PythonでスパイラルマトリックスIIを実装！螺旋状の正方行列を生成するアルゴリズム

スパイラルマトリックスIIとは

正の整数 n が与えられたとき、1 から n² までの数字を外側から内側へ渦巻き状に配置した n×n の正方行列を生成するのが「スパイラルマトリックスII」の問題です。

例えば n = 4 の場合、生成される行列は以下のようになります。

数字が時計回りの渦巻き状に並んでいるのが分かります。この記事では、この行列を Python で生成するアルゴリズムの考え方と実装例をわかりやすく解説します。

アルゴリズムの考え方

基本となるアイデアは、「行列の外周を上→右→下→左の順に埋め、1周終わるごとに境界を内側へ1つずつ狭めていく」というものです。具体的な手順は以下の通りです。

1. (row1, col1) = (0, 0)、(row2, col2) = (n, n) として境界を設定します。全要素が 0 の n×n 行列 res を作成し、num を 1 で初期化します。
2. num ≤ n² の間、以下を繰り返します。
   • 上端の行：col1 から col2−1 まで、左から右へ res[row1][i] = num を代入しながら num を増やします。
   • 右端の列：row1+1 から row2−1 まで、上から下へ res[i][col2−1] = num を代入します。
   • 下端の行：col2−2 から col1 まで、右から左へ res[row2−1][i] = num を代入します。
   • 左端の列：row2−2 から row1+1 まで、下から上へ res[i][col1] = num を代入します。
   • 各方向の処理の後、num が n² を超えていればループを抜けます。
   • 1周埋め終えたら、row1 を +1、row2 を −1、col1 を +1、col2 を −1 として境界を内側へ縮めます。
3. 最終的に res を返します。

Pythonでの実装例

それでは、実際のコードを見てみましょう。

class Solution(object):
    def generateMatrix(self, n):
        row1 = 0
        col1 = 0
        row2 = n
        col2 = n
        result = [[0 for i in range(n)] for j in range(n)]
        num = 1
        while num <= n**2:
            # 上端の行を左から右へ埋める
            for i in range(col1, col2):
                result[row1][i] = num
                num += 1
            if num > n**2:
                break
            # 右端の列を上から下へ埋める
            for i in range(row1+1, row2):
                result[i][col2-1] = num
                num += 1
            if num > n**2:
                break
            # 下端の行を右から左へ埋める
            for i in range(col2-2, col1-1, -1):
                result[row2-1][i] = num
                num += 1
            if num > n**2:
                break
            # 左端の列を下から上へ埋める
            for i in range(row2-2, row1, -1):
                result[i][col1] = num
                num += 1
            # 境界を内側へ1つずつ縮める
            row1 += 1
            row2 -= 1
            col1 += 1
            col2 -= 1
        return result

ob1 = Solution()
print(ob1.generateMatrix(4))

入力

4

出力

[[1, 2, 3, 4], [12, 13, 14, 5], [11, 16, 15, 6], [10, 9, 8, 7]]

計算量について

時間計算量は O(n²) です。行列の全要素をそれぞれちょうど1回ずつ訪問するためです。空間計算量も結果を格納する行列のぶんだけ O(n²) が必要ですが、補助的に使用する追加メモリは O(1) に抑えられています。