ABC410-D を Python で解いた

ABC410 D - XOR Shortest Walk
https://atcoder.jp/contests/abc410/tasks/abc410_d

解けなくて悔しかったので復習のために記事にしてみました。

コンテスト中の自分の動き

C までは楽に終わりましたが D で手が止まりました。いろいろなアイデアが浮かびました。全ての経路を列挙して比較できればいいですが、多重辺はまだしも閉路の扱いが厄介です。最小値を求める際に上の桁から順に桁ごとに考えれば楽になる気がします。前々回 ABC408-E がそんな問題だったと聞きましたし、 XOR を桁ごとにみて 1 の個数の偶奇を調べるのは典型ですね。そもそも経路を全て列挙できるようなら桁ごとに考えるような工夫は不要でしょうし、やはり XOR であることに意味があるのではないかと考えました。というわけで経路を列挙することは諦めました。

その後も考えていくと、桁ごとに考えるのは無理があるという結論に達しました。最上位桁が 0 になるような walk を把握しておき、次はそれらの中から次の桁が 0 になるような walk を探し……みたいになっていくと結局経路を列挙できなければいけません。これは先ほども書いたように多重辺やら閉路やらで扱いきれません。ここでようやく桁ごとに処理する方針を抹消します。

E問題を考えたりしている間に、重みが最大で 1023 までしかないことに気づきます。各頂点ごとに、その頂点に辿り着いたときの XOR の最小値を記録していく DP はどうか？と思いつきました。ただ、最小値だけを記録してもいけません。XOR の値は掛け合わせる数字次第でいくらでも上下してしまうので、取り得る値を全て記録しておく必要があります。

これを実装しました。お察しの通り結果は TLE でした。7つ出ました。以下のコードを見ていただければわかるように、G[v] と dp[v] を掛け合わせる2重ループがあります。dp[v] に入る値の個数が最大 1024 個しかないといってもやはりこんな風にループさせるのは無茶だったようです。結局ここから解決の糸口が見えず、不正解のまま終わってしまいました。

from collections import deque
N, M = map(int, input().split())
dp = [set() for _ in range(N+1)]
dp[1] = {0}
G = [[] for _ in range(N+1)]
for i in range(M):
  a, b, w = map(int, input().split())
  G[a].append((b, w))

que = deque([1])
while(que):
  v = que.popleft()
  for (v2, w) in G[v]: #### このループと
    v_xors = []
    for v_xor in dp[v]: ###### このループが重なるのが嫌な予感
      v_xors.append(v_xor)
    for v_xor in v_xors:
      new_xor = v_xor ^ w
      if new_xor not in dp[v2]:
        dp[v2].add(new_xor)
        que.append(v2)

if len(dp[N]) == 0:
  print(-1)
else:
  print(min(dp[N]))

余談ですが E 問題も解けませんでした。解説を見たら簡単だったのでこれすら思いつけないのかと落ち込みました。DPが弱い自覚はあるので特に鍛えないといけないかもしれません。結局3完に終わり、直近3回のABCでレートを70も落とす大惨事となりました……。

コンテスト後

まず大人しく公式解説を読みました。頂点倍化。初めて聞く単語で、完全に目から鱗でした。ただ頂点を増やすアイデア自体は知っていたはずで、実際に ABC395-E はこれを用いてコンテスト中に正解しています。表の世界と裏の世界を用意して両者を行き来するというやつですね。

ただ、2倍どころか 1024 倍にできるなんて想像もしておらず完全に虚を突かれました。

実装

最初にグラフを作る際に頂点番号と XOR の値をそれぞれ持たせるようにします。辺を張る際、相方になる頂点は頂点番号と XOR の値の組（tuple）で表現します。後で BFS をする際にもこの組を deque に入れておくようにします。あとは基本的な BFS ですね。

from collections import deque
N, M = map(int, input().split())
G = [[[] for _ in range(2**10)] for _ in range(N+1)]
for i in range(M):
  a, b, w = map(int, input().split())
  for bit in range(2**10):
    G[a][bit].append((b, bit ^ w))

seen = [[False for _ in range(2**10)] for _ in range(N+1)]
seen[1][0] = True
que = deque([(1, 0)])
while(que):
  (v, bit) = que.popleft()
  for (v2, bit2) in G[v][bit]:
    if seen[v2][bit2]:
      continue
    seen[v2][bit2] = True
    que.append((v2, bit2))
ans = -1
for bit in range(2**10):
  if seen[N][bit]:
    ans = bit
    break
print(ans)

まとめ

最大 1024 にしかならないことに着目するところまでは良かったようです。あり得る XOR の値を全て持っておくのもあながち間違ってはいなかったと思います。ただそれを表現する方法がまずかったんですね。頂点数を増やすことでとり得る値を全て把握する必要がありました。

ある頂点からの最短距離を求める問題みたいに頂点ごとに 1 種類の値を持たせればいい場合なら基本的な BFS でいいですが、この問題みたいに頂点ごとに複数の値を持たせる必要がある場合は頂点そのものを増やすことで対応できるということで自分の中では納得しました。今後この考え方を柔軟に使えるようにしたいものです。

また、これはDPにも似ている気がしました。あれも

DP[i][j] := i 個目までの荷物から何個か選んで合計の価値を j にできるかどうか（$ i \leq 100 , j \leq 1000 $）

みたいにしますよね。取り得る値の範囲が小さいことを利用して全ての値を網羅できる配列を用意してしまうという豪快な手法です。そもそも BFS 自体が DP に似ている気もします。そんなことを考えていたら少し理解が深まった気がしました。