Pythonで全員が集まる最適な場所の最小移動距離を求めるプログラム

問題概要

2次元の行列（グリッド）が与えられ、そこには以下のような値が含まれています。

• 0：空きセル
• 1：壁
• 2：人

人は上下左右の4方向に移動できます。このとき、壁以外のセルの中から、すべての人がそのセルまで歩く合計距離が最小になる「集合場所」を見つけ、その最小距離を求めるのが目的です。

入力例

この場合の出力は 7 になります。最適な集合場所は右下の角です。

解法のアプローチ

この問題は、各人の位置から幅優先探索（BFS）を実行し、グリッド上のすべてのセルへの移動コストを計算することで解けます。具体的な手順は以下の通りです。

1. twos（人の位置ごとのキュー）と costs（各人から各セルへのコスト表）という2つのマップを用意します。

2. 行列の各セルを走査し、値が 2（人）であれば、その位置を twos に登録し、costs には行列と同じサイズの無限大で初期化した2次元配列を作成します。

3. twos の各エントリに対してBFSを実行します。

   • 訪問済みセルを記録するセット seen を用意します。
   • キューが空になるまで、先頭の要素 (i, j, cost) を取り出します。
   • すでに訪問済みであればスキップし、未訪問なら seen に追加して costs[k][i][j] にコストを記録します。
   • 上下左右の4方向について、行列の範囲内かつ壁でないセルをキューの末尾に cost + 1 として追加します。

4. ans を無限大で初期化し、すべてのセルについて各人からのコストの合計を計算し、その最小値を更新していきます。

5. 最終的な ans を返します。

Pythonでの実装例

class Solution:
    def solve(self, matrix):
        twos = {}
        costs = {}
        for i, r in enumerate(matrix):
            for j, v in enumerate(r):
                if v == 2:
                    twos[(i, j)] = [(i, j, 0)]
                    costs[(i, j)] = [[1e9 for _ in matrix[0]] for _ in matrix]
        for k, q in twos.items():
            seen = set()
            while q:
                i, j, cost = q.pop(0)
                if (i, j) in seen:
                    continue
                seen.add((i, j))
                costs[k][i][j] = cost
                for di, dj in ((1, 0), (-1, 0), (0, 1), (0, -1)):
                    ni, nj = i + di, j + dj
                    if (ni >= 0 and nj >= 0 and ni < len(matrix) and nj < len(matrix[0]) and matrix[ni][nj] != 1):
                        q.append((ni, nj, cost + 1))
        ans = 1e9
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                cur_cost = 0
                for arr in costs.values():
                    cur_cost += arr[i][j]
                ans = min(ans, cur_cost)
        return ans
ob = Solution()
matrix = [
    [2, 0, 1, 0],
    [1, 0, 1, 2],
    [0, 0, 2, 2]
]
print(ob.solve(matrix))

入力

matrix = [
[2, 0, 1, 0],
[1, 0, 1, 2],
[0, 0, 2, 2]]

出力

7

計算量について

各人の位置に対してBFSを1回ずつ実行するため、時間計算量は O(P × R × C)（Pは人数、R×Cはグリッドのサイズ）となります。空間計算量も同様に O(P × R × C) であり、人数やグリッドが大きくなるとメモリ使用量に注意が必要です。また、実装では q.pop(0) を使っていますが、パフォーマンスを重視する場合は collections.deque を使うことでキュー操作を効率化できます。