蟻本初級編攻略 - 2-1 Lake Counting
やはり以前もブログに記事を書いた問題。
やはりコードが長い、というかごちゃごちゃしている・・・
今だったらこう書く:
from itertools import product n, m = map(int, input().split()) lake = [input() for _ in range(n)] lakedict = {(i, j):c for i, row in enumerate(lake) for j, c in enumerate(row)} def dfs(i, j): if lakedict.get((i, j), '.') != 'W': return lakedict[i, j] = '.' for di, dj in product((-1, 0, 1), repeat=2): dfs(i+di, j+dj) n = 0 for i, j in lakedict: if lakedict[i, j] == 'W': n += 1 dfs(i, j) print(n)
ATC001A - 深さ優先探索
失敗1:Recursion Limit
まずは非常に素直にDFSを実装してみた:
from itertools import product h, w = map(int, input().split()) xs = [input() for _ in range(h)] mapdict = {(i, j):c for i, row in enumerate(xs) for j, c in enumerate(row)} def dfs(i, j): current = mapdict.get((i, j), '#') if current == '#': return False if current == 'g': return True mapdict[i, j] = '#' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) return any(dfs(i+di, j+dj) for di, dj in dirs) start = next(coord for coord, c in mapdict.items() if c == 's') print('Yes' if dfs(*start) else 'No')
サンプルケースは全部通るし、コード自体も見通しは悪くないのでバグはなさそう。ということで走らせてみると見事に複数のRE。再帰リミットにぶつかっているようだ。
from sys import setrecursionlimit setrecursionlimit(5000000)
というわけで再帰リミットをあげてみた。これでほとんど通るのだが、ひとつだけMLE。
Pythonでは非常に深い再帰の問題は再帰リミットだけではなく、関数スタックが大きいのでメモリ使用が厳しいという点もある。末尾呼び出し最適化もないし。
というわけで、まあ再帰しなければいいよね。という単純な解決。
from itertools import product h, w = map(int, input().split()) xs = [input() for _ in range(h)] mapdict = {(i, j):c for i, row in enumerate(xs) for j, c in enumerate(row)} def dfs(start): stack = [start] while stack: i, j = stack.pop() current = mapdict.get((i, j), '#') if current == '#': continue if current == 'g': return True mapdict[i, j] = '#' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) stack.extend((i+di, j+dj) for di, dj in dirs) return False start = next(coord for coord, c in mapdict.items() if c == 's') print('Yes' if dfs(start) else 'No')
この問題では可能か否かを葉ノードで判定してそのブール値をそのまま木の上まで返しているだけなので、スタックを使ってしまえばあまりコードをいじらずに再帰を避けることができる。
これは無事にAC。
ARC031B - 埋め立て
似たような問題と似たような解法。
from functools import reduce from operator import and_ mapdict = {(i, j): c for i in range(10) for j, c in enumerate(input())} def dfs(start): stack = [start] while stack: i, j = stack.pop() current = mapdict.get((i, j), 'v') if current == 'v': continue if current == 'x': yield (i, j) continue mapdict[i, j] = 'v' dirs = ((-1, 0), (1, 0), (0, 1), (0, -1)) stack.extend((i+di, j+dj) for di, dj in dirs) surroundings = [] for coord in mapdict: if mapdict[coord] == 'o': surroundings.append(set(dfs(coord))) print('YES' if len(reduce(and_, surroundings)) else 'NO')
yieldを使って島の周囲の海の座標をシーケンスとして返している。それをset化し、「各島の周りの海の座標のset」すべてに共通する要素があるかを判定する。
ARC037B - バウムテスト
無指向グラフにいくつ非巡回グラフが含まれているかを数える問題。
問題文をみると???だが実際に解法は非常に似通っている:
from collections import defaultdict n, m = map(int, input().split()) graph = defaultdict(list) for _ in range(m): u, v = map(int, input().split()) graph[u].append(v) graph[v].append(u) visited = set() def dfs(v): stack = [(None, v)] while stack: last, current = stack.pop() if current in visited: return 0 visited.add(current) stack.extend((current, node) for node in graph[current] if node != last) return 1 ans = 0 for node in range(1, n+1): if node not in visited: ans += dfs(node) print(ans)
無指向グラフでの巡回を探知する時に、「たった今きたノードに戻れる事実は『巡回』として認識しない」というロジックを組み込むのが一番めんどくさかった。スタックに今のノードと次に調べるべきノードの両方をtupleとして乗せることで対応している。