[Python] ABC004 D 動的計画法

問題

D – マーブル

動的計画法を使って解く。

回答

動的計画法

参考

Hacking to the Gate ! 別館　AtCoder Beginer Contest 004 – D

import sys
# input処理を高速化する
input = sys.stdin.readline


def chmin(a, b):
    if a <= b:
        return a
    return b

def main():
    R, G, B = map(int, input().split())

    NUM_MARBLES = R + G + B
    NUM_BOXES = 1000

    # dp[今見ている場所][マーブルの残り数] = 操作回数
    # 最小化問題なのでinfに初期化
    f_inf = float('inf')
    dp = [[f_inf] * (NUM_MARBLES+1) for _ in range(NUM_BOXES)]
    # まだ一度も操作していないときはdp = 0
    for i in range(NUM_BOXES):
        dp[i][NUM_MARBLES] = 0

    # マーブルの残り数に対して移動量の計算が変わる
    # cost(pos, remain)はremain個残っているときにpos番目の箱に入れるための最小の移動数。
    # 左から順に入れているので、残りの数によってどの色を入れるのが移動数最小になるのかは一意に定まり、
    # その色の箱の位置とposの差の絶対値が移動量になる。
    # スタート位置は R(400), G(500), B(600) とする。
    def cost(pos, remain):
        if remain >= G + B:
            return abs(400 - pos)
        elif remain >= B:
            return abs(500 - pos)
        else:
            return abs(600 - pos)

    for i in range(1, NUM_BOXES):
        for j in range(NUM_MARBLES):
            dp[i][j] = chmin(dp[i-1][j], dp[i-1][j + 1] + cost(i, j))

    # マーブルの残り数が0の時のdpの最小値
    ans = f_inf
    for i in range(NUM_BOXES):
        ans = chmin(dp[i][0], ans)
    
    print(ans)


main()