Lake Counting
蟻本&POJから水たまりを数える問題:
ナイーブに書くなら:
import itertools as it def solve(map_): lakes = {(i, j):set((i,j)) for i, row in enumerate(map_) for j, v in enumerate(row) if v == 'W'} def neighbors(i, j): n = ((i+i2, j+j2) for i2, j2 in it.product((-1,0,1), repeat=2)) return [c for c in n if c in lakes] def merge(lake1, lake2): if len(lake1) < len(lake2): lake1, lake2 = lake2, lake1 lake1.update(lake2) for cell in lake2: lakes[cell] = lake1 for cell in lakes: for neighbor in neighbors(*cell): lake1, lake2 = lakes[cell], lakes[neighbor] if lake1 != lake2: merge(lake1, lake2) return len(set(tuple(v) for v in lakes.values()))
merge関数に「小さい方の水たまりのメンバを順次アップデート」という処理がある。ちゃんと考えていないが多分全体でO(nlogn)か。
しかし蟻本に載っている「つながっているWをDFSで消していって、DFSを開始した回数が答え」というのはすごくエレガントだ。
import itertools as it def solve(map_): def dfs(i, j): map_[i][j] = '.' for i1, j1 in it.product((-1,0,1), repeat=2): n, m = i+i1, j+j1 if (n==0 and m==0) or n==-1 or n==len(map_) or m==-1 or m==len(map_[0]): continue if map_[i][j] == 'W': dfs(n, m) count = 0 for i in range(len(map_)): for j in range(len(map_[0])): if map_[i][j] == 'W': count += 1 dfs(i, j) return count
二重ループがcellの総数だけ回って、dfsは'W'の数だけ呼ばれて再帰以外の部分はO(1)なので、地図のdimensionをNM、'W'の数をWで表せば、全体の計算量はO(NM)+O(W)になる。
こういう「探索しているデータ構造を逐次変更していくことで計算量を減らす」という解法はパッと思いつかないので勉強になる。