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)

この記事が気に入ったらサポートをしてみませんか?