Open In Colab

はじめに

このノートブックはEducational DP contestのA~Z問題の解答をPythonで実装したものです。

解答コードはC++言語で書かれたQiita記事「動的計画法超入門」の解答をPythonにほぼそのまま焼き直したものです。考え方はネタ元のQiita記事が凄くわかりやすいのでそちらを参考にしてください。 「Python言語使用者でDPを勉強する際のお役に立てれば」と思い公開しています。

※ コードのアウトプットがコンテストの出力例と合致していることはテスト済み。

In [1]:
import numpy as np
import math

問題A:Frog1

集めるDP

In [ ]:
def frog1_aggr(n, h):
    # DP配列の初期化
    dp = np.full(n, np.inf)

    # 初期条件設定
    dp[0] = 0

    # ループ
    for i in range (1, n):
        dp[i] = min(dp[i], dp[i-1] + abs(h[i] - h[i - 1]))
        if i > 1:
            dp[i] = min(dp[i], dp[i-2] + abs(h[i] - h[i - 2]))

    return dp[-1]
In [ ]:
#### 入力
# n = 4
# h = [10, 30, 40, 20]
# n = 2
# h = [10, 10]
n = 6
h = [30, 10,60, 10, 60, 50]
In [ ]:
frog1_aggr(n, h)
Out[ ]:
40.0

配るDP

In [2]:
def frog1_dist(n, h):
    # DP配列の初期化
    dp = np.full(n, np.inf)

    # 初期条件設定
    dp[0] = 0

    # ループ
    for i in range (0, n):
        target = i + 1
        if(target) < n:
            dp[target] = min(dp[target], dp[i] + abs(h[i] - h[target]))
        target = i + 2
        if(target) < n:
            dp[target] = min(dp[target], dp[i] + abs(h[i] - h[target]))

    return dp[-1]
In [8]:
#### 入力
# n = 4
# h = [10, 30, 40, 20]
# n = 2
# h = [10, 10]
n = 6
h = [30, 10,60, 10, 60, 50]
In [9]:
frog1_dist(n, h)
Out[9]:
40.0

問題B:Frog2

配るDPで実装

In [13]:
def frog2_dist(n, h, k):
    # DP配列の初期化
    dp = np.full(n, np.inf)

    # 初期条件設定
    dp[0] = 0

    # ループ
    for i in range (0, n):
        for m in range(k):
            target = i + m + 1
            if(target) < n:
                dp[target] = min(dp[target], dp[i] + abs(h[i] - h[target]))

    return dp[-1]
In [20]:
#### 入力
# n = 5
# k = 3
# h = [10, 30, 40, 50, 20]
n = 3
k = 1
h = [10, 20, 10]
n = 2
k = 100
h = [10, 10]
n = 10
k = 4
h = [40, 10, 20, 70, 80, 10, 20, 70, 80, 60]
In [21]:
frog2_dist(n, h, k)
Out[21]:
40.0
In [ ]: