Pythonで最大長のスネークシーケンスを見つける方法

スネークシーケンスとは

数値が格納された2次元グリッド(マトリックス)が与えられたとき、その中からスネークシーケンスを見つけて返す問題を考えてみましょう。候補が複数存在する場合は、どれか1つを返せば十分です。

スネークシーケンスとは、グリッド内の隣接するセルの数値をつなげて作る列のことです。各セルにおいて、右隣または下隣のセルの値は、現在のセルの値に対して ±1 でなければなりません。つまり、現在位置がセル (a, b) にある場合、右隣のセル (a, b+1) の値が ±1 なら右へ移動でき、下隣のセル (a+1, b) の値が ±1 なら下へ移動できます。

例として使うグリッド

このグリッドの場合、出力は以下のようになります。

• 最大長:6
• シーケンス:10 (0, 0) → 9 (1, 0) → 8 (1, 1) → 7 (1, 2) → 6 (1, 3) → 7 (2, 3) → 8 (3, 3)

解き方(アルゴリズム)

この問題は動的計画法(DP)を使うことで効率よく解けます。lookup テーブルに「各セルを終点とするスネークシーケンスの最大長」を記録していきます。

パス復元関数 get_path()

• 関数 get_path() を定義します。引数は grid、mat、i、j です。
• path := 空のリストを作成します。
• pt := 点 [i, j] を作成し、path の末尾に追加します。
• grid[i][j] が 0 でない間、次を繰り返します。
  • i > 0 かつ grid[i][j] − 1 == grid[i-1][j] の場合:pt := [i-1, j] を path に追加し、i := i − 1 とします。
  • そうでなく、j > 0 かつ grid[i][j] − 1 == grid[i][j-1] の場合:pt := [i, j-1] を path に追加し、j := j − 1 とします。
• path を返します。

メイン処理

• lookup := M × N サイズの表を作成し、すべて 0 で初期化します。
• length_max := 0、max_row := 0、max_col := 0 とします。
• i を 0 から M-1 まで、j を 0 から N-1 まで走査します(i または j が非ゼロの場合のみ処理)。
  • i > 0 かつ |grid[i-1][j] − grid[i][j]| == 1 の場合:lookup[i][j] = max(lookup[i][j], lookup[i-1][j] + 1) と更新し、length_max より大きければ length_max・max_row・max_col を更新します。
  • j > 0 かつ |grid[i][j-1] − grid[i][j]| == 1 の場合:同様に lookup[i][j] = max(lookup[i][j], lookup[i][j-1] + 1) と更新し、必要なら最大値情報を更新します。
• length_max を表示します。
• path := get_path(lookup, grid, max_row, max_col) でパスを復元し、逆順に出力します。

実装例

それでは、理解を深めるために実際の実装を見てみましょう。

M = 4
N = 4

def get_path(grid, mat, i, j):
    path = list()
    pt = [i, j]
    path.append(pt)
    while (grid[i][j] != 0):
        if (i > 0 and grid[i][j]-1 == grid[i-1][j]):
            pt = [i-1, j]
            path.append(pt)
            i -= 1
        elif (j > 0 and grid[i][j]-1 == grid[i][j-1]):
            pt = [i, j-1]
            path.append(pt)
            j -= 1
    return path

def get_sequence(grid):
    lookup = [[0 for i in range(N)] for j in range(M)]
    length_max = 0
    max_row = 0
    max_col = 0
    for i in range(M):
        for j in range(N):
            if (i or j):
                if (i > 0 and abs(grid[i-1][j] - grid[i][j]) == 1):
                    lookup[i][j] = max(lookup[i][j], lookup[i-1][j] + 1)
                    if (length_max < lookup[i][j]):
                        length_max = lookup[i][j]
                        max_row = i
                        max_col = j
                if (j > 0 and abs(grid[i][j-1] - grid[i][j]) == 1):
                    lookup[i][j] = max(lookup[i][j], lookup[i][j-1] + 1)
                    if (length_max < lookup[i][j]):
                        length_max = lookup[i][j]
                        max_row = i
                        max_col = j
    print("Maximum length:", length_max)
    path = get_path(lookup, grid, max_row, max_col)
    print("Sequence is:")
    for ele in reversed(path):
        print(grid[ele[0]][ele[1]], " [", ele[0], ", ", ele[1], "]", sep="")

grid = [
    [10, 7, 6, 3],
    [9, 8, 7, 6],
    [8, 4, 2, 7],
    [2, 2, 2, 8]]
get_sequence(grid)

入力

[[10, 7, 6, 3],
[9, 8, 7, 6],
[8, 4, 2, 7],
[2, 2, 2, 8]]

出力

Maximum length: 6
Sequence is:
10 [0, 0]
9 [1, 0]
8 [1, 1]
7 [1, 2]
6 [1, 3]
7 [2, 3]
8 [3, 3]

時間計算量は O(M×N)、空間計算量も O(M×N) であり、全セルを一度だけ走査して DP 表を埋めるため、大きなグリッドでも効率的に動作します。