947. Most Stones Removed with Same Row or Column

Just a moment...

二次元座標平面上において、x座標またはy座標が共通の点を次々に削除していき、削除した個数を求める問題。

深さ優先探索

全ての座標から出発し、DFSを行う。

たどり着けた点を削除し、もう削除できなくなったら次のスタート地点から再度探索を行う。

残った点が削除できない点群なので、最初の個数から引いたものが答え。

まずは全座標を辞書に格納。

coord = {(i, j) for i, j in stones}

探索用に同じx座標とy座標の点を列挙。

from collections import defaultdict

row = defaultdict(list)
col = defaultdict(list)

for i, j in stones:
    row[i].append(j)
    col[j].append(i)

DFSを行う。

たどり着いた点を削除し、x方向またはy方向に進んでいく。

辞書から値を削除する際、含まれていない値をremoveするとエラーが出てしまう。

discardを使うと存在しない値はスルーしてくれるので書くのが楽ちん。

def dfs(x, y):
    coord.discard((x, y))
    for i in col[y]:
        if (i, y) in coord:
            dfs(i, y)
    for i in row[x]:
        if (x, i) in coord:
            dfs(x, i)

あとは全ての座標をスタート地点として全探索すれば完了。

ans = len(stones)
for i, j in stones:
    if (i, j) in coord:
        dfs(i, j)
        ans -= 1

return ans

全部繋げるとこちら。

class Solution:
    def removeStones(self, stones: List[List[int]]) -> int:
        from collections import defaultdict

        coord = {(i, j) for i, j in stones}
        row = defaultdict(list)
        col = defaultdict(list)

        for i, j in stones:
            row[i].append(j)
            col[j].append(i)

        def dfs(x, y):
            coord.discard((x, y))
            for i in col[y]:
                if (i, y) in coord:
                    dfs(i, y)
            for i in row[x]:
                if (x, i) in coord:
                    dfs(x, i)

        ans = len(stones)
        for i, j in stones:
            if (i, j) in coord:
                dfs(i, j)
                ans -= 1

        return ans
            

UnionFind

DFSでは座標を削除していったが、発想を変えると「あるスタート地点から辿れる点群の集合の個数」とも捉えることが出来るため、集合を扱うデータ構造のUnionFindが使える。

今回は各集合に含まれる個数は考えなくていいため、少し実装が楽になる。

まずはUnionFindのクラスを実装する。

今回の問題は各集合に含まれる個数は考えなくていいため、少し実装が楽になる。

集合の親となるインデックスを格納する配列を初期化。

-1が集合の親で、正の値は自分の親のインデックスを表す。

class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * n

集合の親は再帰的に探索する。

探索途中で親を変更しておくと定数倍高速化できる。

    def find(self, x):
        if self.parent[x] == -1:
            return x

        self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

2つの値を入力し、片方の親をもう片方の親に書き換えてマージしていく。

親が同じ場合は特に何もしない。

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        self.parent[x] = y

以上でUnionFindが完成。

class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * n

    def find(self, x):
        if self.parent[x] == -1:
            return x

        self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        self.parent[x] = y

与えられたデータからUnionFindを構築していく。

入力は二次元座標のリストで与えられているが、何番目と何番目の座標が同じ集合なのかを求めたいので、インデックスで集合を構築していく。

n = len(stones)

uf = UnionFind(n)

for i in range(n):
    for j in range(i + 1, n):
        x1, y1 = stones[i]
        x2, y2 = stones[j]
        if x1 == x2 or y1 == y2:
            uf.union(i, j)

集合の親を探索していく。

親は削除されない点を意味するので、最初の個数から残った親の個数を引けば答えになる。

ans = n
for i in uf.parent:
    if i == -1:
        ans -= 1

return ans

全部繋げて提出。

class Solution:
    def removeStones(self, stones: List[List[int]]) -> int:
        n = len(stones)

        uf = UnionFind(n)

        for i in range(n):
            for j in range(i + 1, n):
                x1, y1 = stones[i]
                x2, y2 = stones[j]
                if x1 == x2 or y1 == y2:
                    uf.union(i, j)

        ans = n
        for i in uf.parent:
            if i == -1:
                ans -= 1

        return ans


class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * n

    def find(self, x):
        if self.parent[x] == -1:
            return x

        self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        self.parent[x] = y
            
        

コメント

タイトルとURLをコピーしました