Pythonのheapqとdeque
AtCoder ABC088 D Grid Repainting
台風で外出できない連休だったので,家でAtCoderをやっていた。今回はこれ。
問題文を読んだ瞬間に,「あ,BFSだな」というのはわかりました。蟻本を読んでいましたからね!で,見様見真似で実装したコードがこちら
import sys
import os
import heapq
INF = 10**5
H, W = map(int,input().split())
S = ["" for _ in range(H)]
cntSharp = 0
for i in range(H):
S[i] = input()
cntSharp += S[i].count("#")
dp = [[INF for _ in range(W)] for _ in range(H)]
dp[0][0] = 0
#移動用のベクトル
dx = [1,0,-1,0]
dy = [0,1,0,-1]
P = [(0,0)]
heapq.heapify(P)
while P:
p = heapq.heappop(P)
s,t = p[0], p[1]
for i in range(4):
nx,ny = s + dx[i], t + dy[i]
if 0 <= nx and 0 <= ny and nx <= H-1 and ny <= W-1 and S[nx][ny] =="." and dp[nx][ny] == INF:
dp[nx][ny] = dp[s][t] + 1
heapq.heappush(P,(nx,ny))
if dp[H-1][W-1] == INF:
print(-1)
else:
print(H*W - dp[H-1][W-1] - cntSharp - 1)
このコードで15/20がAC。残り5ケースがWAとなりました。正直全然わからず,公式editorial見ても有益な情報は得られませんでした。そこで,AtCoderに提出されている他の方のコードを見ると,heapqでなくdequeを使っていました。正直違いはそれくらいでした。
dequeってなんだ?
ちなみにheapqはこちら
dequeはappendで右側に追加できる。popleftで左から取り出せる。heapqは常に最小値を取り出す。今回はqueにtuple型で値を格納しています。おそらくこんなことが起こっていたのでしょう。
A = [(1,0),(0,1),(-1,0),(0,-1)]
P = []
heapq.heapify(P)
Q = deque([])
for a in A:
heapq.heappush(P,a)
Q.append(a)
while P:
print("P:",heapq.heappop(P))
print()
while Q:
print("Q:",Q.popleft())
実行結果
P: (-1, 0)
P: (0, -1)
P: (0, 1)
P: (1, 0)
Q: (1, 0)
Q: (0, 1)
Q: (-1, 0)
Q: (0, -1)
Pの出力とQの出力の順番が異なっています。
heapqだとうまくいかなかった理由
今回のBFSの処理の中で以下のような処理があります。
dx = [1,0,-1,0]
dy = [0,1,0,-1]
for i in range(4):
nx,ny = s + dx[i], t + dy[i]
if 0 <= nx and 0 <= ny and nx <= H-1 and ny <= W-1 and S[nx][ny] =="." and dp[nx][ny] == INF:
dp[nx][ny] = dp[s][t] + 1
heapq.heappush(P,(nx,ny))
現在地(nx, ny)に隣接する4方向を順番に調べていますが,その順番は事前に定義したdx, dyという2つのlistに依存しています。以下のような順番で走査しています。
(nx+1,ny) -> (nx,ny+1) -> (nx-1,ny) -> (nx, ny-1)
これで,この処理においてdequeを使用する場合とheapqを使用する場合とで違いがあることがわかりました。今回実装するBFSにおいては例えば
1.(3, 3)を調べる
2.(2, 3), (3, 2), (3, 4), (4, 3)の4か所を調べる
3.(2, 3)の周囲4か所を調べる★,(3, 2)の周囲4か所を調べる★★
(3, 4)の周囲4か所を調べる☆(4, 3)の周囲4か所を調べる☆☆
ということを順に行います。この<4か所を調べる>ことはひと塊の処理であり,調べ終わる前に別の4か所に移ってはいけません(★が終わってから★★へ移るというようにしないといけない)。しかし,heqpqを使用した場合は,★,★★,☆,☆☆に含まれる16か所の中から小さい順に調べてしまいます。そうすると遠回りしてしまうマス目が出てしまいます。ということでheapqの部分をdequeに変更して(2,3)の周囲4マスをappend,(3,4)の周囲4マスをappendという風にしておけば,4個の塊を崩すことなく,探索を行うことができます。
ACだった実装
import sys
import os
from collections import deque
INF = 10**5
H, W = map(int,input().split())
S = ["" for _ in range(H)]
cntSharp = 0
for i in range(H):
S[i] = input()
cntSharp += S[i].count("#")
dp = [[INF for _ in range(W)] for _ in range(H)]
dp[0][0] = 0
#移動用のベクトル
dx = [1,0,-1,0]
dy = [0,1,0,-1]
P = deque([])
P.append((0,0))
while P:
p = P.popleft()
s,t = p[0], p[1]
for i in range(4):
nx,ny = s + dx[i], t + dy[i]
if 0 <= nx and 0 <= ny and nx <= H-1 and ny <= W-1 and S[nx][ny] =="." and dp[nx][ny] == INF:
dp[nx][ny] = dp[s][t] + 1
P.append((nx,ny))
if dp[H-1][W-1] == INF:
print(-1)
else:
print(H*W - dp[H-1][W-1] - cntSharp - 1)