Lake Counting

蟻本&POJから水たまりを数える問題：

2386 -- Lake Counting

ナイーブに書くなら：

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)になる。

こういう「探索しているデータ構造を逐次変更していくことで計算量を減らす」という解法はパッと思いつかないので勉強になる。