【Python】グリッドに陸地ブロックを1つずつ追加しながら島の数を求めるプログラム

無限に広がる水のグリッドがあるとします。このグリッドに対して、陸地のブロックを1つずつ追加していきます。座標のリスト land_requests が与えられ、各座標は [r, c] の形式で表されます(r は行、c は列を意味します)。求めたいのは、各ブロックを追加した直後の時点で存在する島の数を要素とするリストです。

例えば、入力が land_requests = [[1, 1], [2, 4], [1, 2], [1, 4], [1, 3]] の場合、出力は [1, 2, 2, 2, 1] となります。

なぜ [1, 2, 2, 2, 1] になるのか

• [1, 1] を追加 → 島は1つ
• [2, 4] を追加 → 既存の島と隣接していないため、島は2つ
• [1, 2] を追加 → [1, 1] と隣接して1つの島に統合されるため、島は2つ
• [1, 4] を追加 → [2, 4] と隣接して統合されるため、島は2つ
• [1, 3] を追加 → 左右の両方の島をつなぐため、すべてが1つになり、島は1つ

解決のアプローチ:Union-Find(素集合データ構造)

この問題は「Union-Find」を使うことで効率的に解けます。新しい陸地が既存の島と隣接するたびにグループを統合し、連結成分(島)の数を管理するイメージです。手順は以下の通りです。

• d := 移動方向のリスト [(-1, 0), (0, 1), (1, 0), (0, -1)] を定義
• idx := 0(各ブロックに振るID)
• mp := 座標とIDを対応付ける辞書
• p := Union-Find 用の親ノードリスト
• size := 各連結成分のサイズを記録するリスト
• comp := 現在の島(連結成分)の数
• ans := 結果を格納するリスト

search() 関数(Find 操作)

引数 u を受け取り、経路圧縮を行いながら根を返します。

• u が p[u] と等しい場合は u を返す
• そうでなければ p[u] := search(p[u]) として経路を圧縮し、p[u] を返す

connect() 関数(Union 操作)

引数 u, v を受け取り、2つのブロックを同じ島に統合します。

• pu := search(u)、pv := search(v) でそれぞれの根を取得
• pu と pv が等しい場合(既に同じ島)は何もしない
• comp を 1 減らす
• size[pu] >= size[pv] の場合:p[pv] := pu とし、size[pu] += size[pv]
• それ以外の場合:p[pu] := pv とし、size[pv] += size[pu]

メイン処理

• land_requests の各 request について以下を実行:
• (i, j) := request として座標を取り出す
• mp[(i, j)] := idx で座標にIDを割り当てる
• p に idx を追加し、size に 1 を追加する
• idx := idx + 1、comp := comp + 1 と更新
• 方向 d の各 k について:ni := i + k[1]、nj := j + k[0] を計算し、(ni, nj) が mp に存在すれば connect(mp[(i, j)], mp[(ni, nj)]) を呼び出して隣接する陸地と統合する
• ans に comp を追加する

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

実装例(Python)

d = [(-1, 0), (0, 1), (1, 0), (0, -1)]

class Solution:
    def search(self, u):
        if u == self.p[u]:
            return u
        self.p[u] = self.search(self.p[u])
        return self.p[u]

    def connect(self, u, v):
        pu = self.search(u)
        pv = self.search(v)
        if pu == pv:
            return
        self.comp -= 1
        if self.size[pu] >= self.size[pv]:
            self.p[pv] = pu
            self.size[pu] += self.size[pv]
        else:
            self.p[pu] = pv
            self.size[pv] += self.size[pu]

    def solve(self, land_requests):
        self.idx = 0
        self.mp = dict()
        self.p = []
        self.size = []
        self.comp = 0
        ans = []
        for request in land_requests:
            i, j = request
            self.mp[(i, j)] = self.idx
            self.p.append(self.idx)
            self.size.append(1)
            self.idx += 1
            self.comp += 1
            for k in d:
                ni = i + k[1]
                nj = j + k[0]
                if (ni, nj) in self.mp:
                    self.connect(self.mp[(i, j)], self.mp[(ni, nj)])
            ans.append(self.comp)
        return ans

ob = Solution()
land_requests = [[1, 1], [2, 4], [1, 2], [1, 4], [1, 3]]
print(ob.solve(land_requests))

入力

[[1, 1],[2, 4],[1, 2],[1, 4],[1, 3]]

出力

[1, 2, 2, 2, 1]