Pythonで各セルに最も近い0までのマンハッタン距離を格納した行列を生成するプログラム

問題の概要

0と1だけで構成される2値行列を考えます。この行列と同じサイズの新しい行列を作成してください。ただし、新しい行列の各セルには、元の行列においてその位置から最も近い0までのマンハッタン距離を格納します。なお、元の行列には少なくとも1つの0が存在すると仮定できます。

たとえば、入力が次のような行列だった場合を考えてみましょう。

このとき、出力は次のようになります。

これは、左下のセルだけが最も近い0までの距離が2になるためです。

解決のための手順

この問題は、動的計画法(DP)の考え方を応用し、行列を2回走査するだけで効率的に解くことができます。具体的な手順は以下の通りです。

• m := 行列の行数、n := 行列の列数 とします。
• y を 0 から m-1 まで繰り返します。
  • x を 0 から n-1 まで繰り返します。
    • matrix[y, x] が 0 以外の場合、matrix[y, x] := 無限大 と初期化します。
• y を 0 から m-1 まで繰り返します(左上からの順方向パス)。
  • x を 0 から n-1 まで繰り返します。
    • y が 0 以外なら、matrix[y, x] = min(matrix[y, x], matrix[y - 1, x] + 1) と更新します。
    • x が 0 以外なら、matrix[y, x] = min(matrix[y, x], matrix[y, x - 1] + 1) と更新します。
• y を m - 1 から 0 まで減らしながら繰り返します(右下からの逆方向パス)。
  • x を n - 1 から 0 まで減らしながら繰り返します。
    • y + 1 < m の場合、matrix[y, x] = min(matrix[y, x], matrix[y + 1, x] + 1) と更新します。
    • x + 1 < n の場合、matrix[y, x] = min(matrix[y, x], matrix[y, x + 1] + 1) と更新します。
• matrix を返します。

アルゴリズムのポイント

まず、0以外のセルをすべて無限大で初期化します。その後、左上から右下へ向かう順方向のパスで「上」と「左」の隣接セルからの距離を伝播させ、続いて右下から左上へ向かう逆方向のパスで「下」と「右」の隣接セルからの距離を伝播させます。この2回の走査により、すべてのセルについて上下左右4方向からの最小距離が正しく求まり、計算量は O(m×n) に抑えられます。

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

サンプルコード(Python)

import math

class Solution:
   def solve(self, matrix):
      m, n = len(matrix), len(matrix[0])
      for y in range(m):
         for x in range(n):
            if matrix[y][x]:
               matrix[y][x] = math.inf
      for y in range(m):
         for x in range(n):
            if y:
               matrix[y][x] = min(matrix[y][x], matrix[y - 1][x] + 1)
            if x:
               matrix[y][x] = min(matrix[y][x], matrix[y][x - 1] + 1)
      for y in range(m - 1, -1, -1):
         for x in range(n - 1, -1, -1):
            if y + 1 < m:
               matrix[y][x] = min(matrix[y][x], matrix[y + 1][x] + 1)
            if x + 1 < n:
               matrix[y][x] = min(matrix[y][x], matrix[y][x + 1] + 1)
      return matrix

ob = Solution()
matrix = [ [1, 0, 1], [1, 0, 1], [1, 1, 0] ]
print(ob.solve(matrix))

入力

[[1, 0, 1],
[1, 0, 1],
[1, 1, 0]]

出力

[[1, 0, 1], [1, 0, 1], [2, 1, 0]]