ABC197 E Traveler

内容

• Atcoder Beginner contest E問題 Travelerをpython でAC

  問題

atcoder.jp

考えたこと

• 各色ごとに、左端で終わるか、右端で終わるかのどちらかが最適。
• 色ごとに試すと、2の(色の数)常になって辛い。
  • 前の色の(左端、右端)だけ見れればいいので、[色数][left or right]のDPにできる。

ポイント

• 最初と最後に番兵を入れると実装が楽

  実装

from collections import defaultdict
N = int(input())

# (pos,color)
D = defaultdict(list)

# colors
Colors = set()
for _ in range(N):
    x, c = map(int, input().split())
    D[c].append(x)
    Colors.add(c)

# 0,max+1は番兵
Colors.add(0)
Colors.add(max(Colors)+1)
Colors = list(sorted(Colors))

# 色ごとの、左端と右端
pos_edge = dict()
# 番兵
pos_edge[0] = (0, 0)
pos_edge[max(Colors)] = (0, 0)
for c, vals in D.items():
    pos_edge[c] = (min(vals), max(vals))


DP = defaultdict(lambda: [float("inf"), float("inf")])
# DP[c][0]:cを見終わったときに　cの左端にいる
# DP[c][1]:cを見終わったときに　cの右端にいる

# スタートは、cost0
DP[0][0] = 0
DP[0][1] = 0

for i in range(1, len(Colors)):
    prev = Colors[i-1]
    cur = Colors[i]

    prev_left, prev_right = pos_edge[prev]
    cur_left, cur_right = pos_edge[cur]

    # (left,right)->left
    for key, p in zip([0, 1], [prev_left, prev_right]):
        # p->right->left
        cost = abs(cur_right-p)+(cur_right-cur_left)
        DP[cur][0] = min(DP[cur][0], DP[prev][key]+cost)

    # (left,right)=>right
    for key, p in zip([0, 1], [prev_left, prev_right]):
        # p->left->right
        cost = abs(cur_left-p)+(cur_right-cur_left)
        DP[cur][1] = min(DP[cur][1], DP[prev][key]+cost)

print(min(DP[Colors[-1]]))

終わりに