🚲

[ABC340 D] Super Takahashi Bros.

2024/02/12に公開

https://atcoder.jp/contests/abc340/tasks/abc340_d

ダイクストラ法に気づけるか問題。

解き方

各頂点から「A と B の 2 つの辺が出ている」と考えることができればもう解けたようなもの。

この例1のグラフでは、なんとなく 1 → 2 → 3 → 5 の順に進むのがよさそう。そのような考え方にはダイクストラ法が適している。

ところが盛大にはまる

普通の競プロ勢はよっぽどのことがない限りハッシュデータ構造なんか使わないが、自分はなんでもハッシュ病に罹っているので、グラフのデータもハッシュで持っていた。それが災いした。

G = Array.new(N) { {} }
ABX.each.with_index do |(a, b, x), i|
  G[i][i.next] = a
  G[i][x.pred] = b
end
G  # => [{1=>100, 2=>200}, {2=>50, 0=>10}, {3=>100, 4=>200}, {4=>150, 1=>1}, {}]

TLE 対策で外側は配列にしたものの「辺とコスト」たちはハッシュで持っていて、それは遅い以外になんら問題はないだろうと考えていた。しかし、TLE ではなく WA になるのは想定外だった。

そこで WA になるテストケースをどうにかして入手したかったのだが、AtCoder が用意してくれている Dropbox には 5 つぐらい前より古いコンテストのテストケースしかなかった。つまり一ヶ月ぐらい待たないと WA になるテストケースがわからないということで、途方に暮れていたのだけど突然閃いた。

  • A → いま 1 面にいて次の面である 2 面に進めるコストは 1
  • B → いま 1 面にいて 2 面にジャンプできるコストは 2

となっていたときにハッシュキーが同じだから A の辺の情報を B の辺の情報で上書きしてしまう。本来ならコスト 1 で進めたはずなのにコスト 2 になってしまう。つまり元のグラフが正しく作れていなかった。

対策として次のように A と B で進む面が同じだった場合にあらかじめコストが低い方を優先してグラフを作る、という工夫が考えらえる。

G = Array.new(N) { {} }
ABX.each.with_index do |(a, b, x), i|
  if i.next == x.pred
    if a < b
      G[i][i.next] = a
    else
      G[i][x.pred] = b
    end
  else
    G[i][i.next] = a
    G[i][x.pred] = b
  end
end
G

が、どっちを優先すればいいかという部分こそダイクストラ先生がやってくれるんだから、自力でわけのわからないことをせずに A と B の2通りを素直に配列で持っとけばいい。

G = Array.new(N) { [] }
ABX.each.with_index do |(a, b, x), i|
  G[i] << [i.next, a]
  G[i] << [x.pred, b]
end
G  # => [[[1, 100], [2, 200]], [[2, 50], [0, 10]], [[3, 100], [4, 200]], [[4, 150], [1, 1]], []]

教訓: 多重辺に気をつけろ

不要な Set を使うべからず

また、調子に乗って配列を Set なんかにしてはいけない。

G = Array.new(N) { Set[] }

これでも通ったけど重複チェックがあるせいか 300 ms ほど遅くなり 1961 ms の TLE 寸前になってしまう。

コード (AC)

N   = gets.to_i                                            # => 5
ABX = N.pred.times.collect { gets.split.collect(&:to_i) }  # => [[100, 200, 3], [50, 10, 1], [100, 200, 5], [150, 1, 2]]

START = 0                                                  # => 0
GOAL  = N.pred                                             # => 4

G = Array.new(N) { [] }
ABX.each.with_index do |(a, b, x), i|
  G[i] << [i.next, a]
  G[i] << [x.pred, b]
end
G                                                          # => [[[1, 100], [2, 200]], [[2, 50], [0, 10]], [[3, 100], [4, 200]], [[4, 150], [1, 1]], []]

d = Array.new(N, Float::INFINITY) # distance
heap = Heap.new { |e| d[e] }

d[START] = 0
heap << START

until heap.empty?
  from = heap.pop
  G[from].each do |to, cost|
    v = d[from] + cost
    if v < d[to]
      d[to] = v
      heap << to
    end
  end
end

p d[GOAL]                                                  # => 350

Discussion