Arantium Maestum

プログラミング、囲碁、読書の話題

蟻本初級編攻略 - 2-1 Lake Counting

やはり以前もブログに記事を書いた問題。

zehnpaard.hatenablog.com

やはりコードが長い、というかごちゃごちゃしている・・・

今だったらこう書く:

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 - 深さ優先探索

atc001.contest.atcoder.jp

失敗1:Recursion Limit

atc001.contest.atcoder.jp

まずは非常に素直に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。再帰リミットにぶつかっているようだ。

atc001.contest.atcoder.jp

from sys import setrecursionlimit
setrecursionlimit(5000000)

というわけで再帰リミットをあげてみた。これでほとんど通るのだが、ひとつだけMLE。

Pythonでは非常に深い再帰の問題は再帰リミットだけではなく、関数スタックが大きいのでメモリ使用が厳しいという点もある。末尾呼び出し最適化もないし。

atc001.contest.atcoder.jp

というわけで、まあ再帰しなければいいよね。という単純な解決。

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 - 埋め立て

arc031.contest.atcoder.jp

arc031.contest.atcoder.jp

似たような問題と似たような解法。

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 - バウムテスト

arc037.contest.atcoder.jp

arc037.contest.atcoder.jp

無指向グラフにいくつ非巡回グラフが含まれているかを数える問題。

問題文をみると???だが実際に解法は非常に似通っている:

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として乗せることで対応している。