[鉄則A17] Dungeon 2

https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_q

A16 - Dungeon 1 の続きで経路も一緒に求める問題。

問題を要約する

• A には順番に移動するときのコストが並んでいる
• B には1つ飛ばしでジャンプするときのコストが並んでいる
• ゴールまでの最小コストとその経路を求めよ

例1のシミュレーション

スタート地点に近い方から最小コストと経路を順に求めていくと、

地点	理由	コスト	経路
1	1からスタートするので	0	1
2	1からのみ移動できるので	2	1-2
3	1から飛んだ方が速いので	5	1-3
4	2から飛んだ方が速いので	5	1-2-4
5	3, 4 どちらからも同じ	8	1-2-4-5 または 1-3-5

となり 1-2-4-5 または 1-3-5 ルートが最小コスト 8 になる。

自力で実装 (TLE)

1116 ms

N = gets.to_i                   # => 5
A = gets.split.collect(&:to_i)  # => [2, 4, 1, 3]
B = gets.split.collect(&:to_i)  # => [5, 3, 3]
START = 0
GOAL  = N.pred

A = [0, *A]                     # => [0, 2, 4, 1, 3]
B = A.take(2) + B               # => [0, 2, 5, 3, 3]

dp = Hash.new(0)
dp[0] = A[0]
dp[1] = A[1]

# 足跡用
predecessor = {}
predecessor[START] = nil
predecessor[1] = START

(2...N).each do |i|
  a = dp[i - 1] + A[i]          # => 6, 6, 8
  b = dp[i - 2] + B[i]          # => 5, 5, 8
  dp[i] = [a, b].min            # => 5, 5, 8
  # どこから来たのか足跡を記録しておく
  if dp[i] == a
    predecessor[i] = i - 1
  else
    predecessor[i] = i - 2
  end
end

# 復元
route_to = -> v { v && [*route_to[predecessor[v]], v.next] }
routes = route_to[GOAL]         # => [1, 2, 4, 5]
p routes.size                   # => 4
puts routes * " "

ダイクストラ法の要領で predecessor に足跡を残していく方法で実装したが TLE になる。随所で連想配列を使っているのが原因と考えて配列にしても効果がない。

パフォーマンス改善後 (AC)

原因はここで、

route_to = -> v { v && [*route_to[predecessor[v]], v.next] }
routes = route_to[GOAL]  # => [1, 2, 4, 5]

これに直したら、

routes = []
from = GOAL
while from
  routes.unshift(from.next)
  from = predecessor[from]
end
routes  # => [1, 2, 4, 5]

通るようになった。再帰で配列要素を展開しまくるのはシンプルに書ける反面、競プロには向いていないのがよくわかった。

全体のコード (120 ms)

N = gets.to_i                   # => 5
A = gets.split.collect(&:to_i)  # => [2, 4, 1, 3]
B = gets.split.collect(&:to_i)  # => [5, 3, 3]
START = 0
GOAL  = N.pred

A = [0, *A]                     # => [0, 2, 4, 1, 3]
B = A.take(2) + B               # => [0, 2, 5, 3, 3]

dp = Hash.new(0)
dp[0] = A[0]
dp[1] = A[1]

predecessor = {}
predecessor[START] = nil
predecessor[1] = START

(2...N).each do |i|
  a = dp[i - 1] + A[i]          # => 6, 6, 8
  b = dp[i - 2] + B[i]          # => 5, 5, 8
  dp[i] = [a, b].min            # => 5, 5, 8
  if dp[i] == a
    predecessor[i] = i - 1
  else
    predecessor[i] = i - 2
  end
end

routes = []
from = GOAL
while from
  routes.unshift(from.next)
  from = predecessor[from]      # => 3, 1, 0, nil
end
p routes.size                   # => 4
routes                          # => [1, 2, 4, 5]
puts routes * " "

想定解の経路復元方法

routes = []
to = GOAL
loop do
  routes.unshift(to.next)
  if to == START
    break
  end
  if dp[to - 1] + A[to] == dp[to]
    to = to - 1
  else
    to = to - 2
  end
end
routes  # => [1, 2, 4, 5]

解説では足跡を記録せずに逆算していた。

全体のコード (115 ms)

N = gets.to_i                   # => 5
A = gets.split.collect(&:to_i)  # => [2, 4, 1, 3]
B = gets.split.collect(&:to_i)  # => [5, 3, 3]
START = 0
GOAL  = N.pred

A = [0, *A]                     # => [0, 2, 4, 1, 3]
B = A.take(2) + B               # => [0, 2, 5, 3, 3]

dp = Hash.new(0)
dp[0] = A[0]
dp[1] = A[1]

(2...N).each do |i|
  a = dp[i - 1] + A[i]          # => 6, 6, 8
  b = dp[i - 2] + B[i]          # => 5, 5, 8
  dp[i] = [a, b].min            # => 5, 5, 8
end

routes = []
to = GOAL
loop do
  routes.unshift(to.next)
  if to == START
    break
  end
  if dp[to - 1] + A[to] == dp[to]
    to = to - 1
  else
    to = to - 2
  end
end
p routes.size                   # => 4
routes                          # => [1, 2, 4, 5]
puts routes * " "