Pythonでグリッド内のボールの着地位置を求めるプログラム

問題概要

m × n のグリッドボックスを考えます。各セルには、左上から右下、もしくは右上から左下へ向けて斜めの板が設置されています。グリッドの上端からボールを落とし、それぞれのボールが底まで到達できるか、そしてどの列に着地するのかを求めるのがこの問題です。

グリッドは行列として与えられ、各セルの値は板の向きを表します。

• 1: 左上から右下へ下る斜めの板
• -1: 右上から左下へ下る斜めの板

n 個のボールを上端の各列から順に落としたとき、底に到達したボールの着地列を答えとして返します。途中で側面の壁に当たったり、V字型の溝にはまって動けなくなったボールについては -1 を出力します。

3 × 3 のグリッドボックスの例

入力例

この行列を入力とした場合、出力は [-1, -1, -1, -1] となります。すべてのボールが途中で引っかかり、底まで到達できないためです。

解法のアプローチ

各ボールについて、グリッドを1行ずつシミュレーションしていきます。具体的な手順は以下の通りです。

• i := mat の行数、j := mat の列数とする
• 結果を格納するための新しいリスト res を用意する
• 開始列 val を 0 から j-1 まで順に処理する:
  • x := val と初期化する
  • r を 0 から i-1 まで繰り返す:
    • s := mat[r][x](現在いるセルの板の向き)
    • x := x + mat[r][x](板に沿って横方向へ移動する)
    • x が範囲外(x < 0 または x ≥ j)になった場合、あるいは移動先のセル mat[r][x] の板の向きが s と異なる場合(V字にはまった状態)は、res の末尾に -1 を追加してループを抜ける
  • 最後の行まで抜けられた場合は、res の末尾に x を追加する
• res を返す

実装上のポイントは、失敗パターンを2種類とも正しく検出することです。ひとつはグリッドの外側の壁に当たるケース、もうひとつは隣接するセルの板が互いに内側を向いており、同じ行の中で左右に振動してしまうケース(いわゆるV字トラップ)です。

Pythonでの実装例

def solve(mat):
    i, j = map(len, (mat, mat[0]))
    res = []
    for val in range(j):
        x = val
        for r in range(i):
            s = mat[r][x]
            x += mat[r][x]
            if x < 0 or x >= j or mat[r][x] != s:
                res += [-1]
                break
        else:
            res += [x]
    return res

print(solve([[1, 1, 1, -1], [-1, 1, 1, -1], [1, -1, -1, 1], [1, -1, 1, -1]]))

入力

[[1, 1, 1, -1], [-1, 1, 1, -1], [1, -1, -1, 1], [1, -1, 1, -1]]

出力

[-1, -1, -1, -1]

計算量の目安

ボール1本につき最大 m 行のシミュレーションを行い、それを n 本分繰り返すため、時間計算量は O(m × n) となります。結果の保存には O(n) の追加メモリを使用します。