二次元座標平面上において、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
コメント