Typical DP Contest A

DP訓練。

tdpc.contest.atcoder.jp

まずはDP抜きにして問題文がきちんと理解できていることを確認するために さくっと実装。

def solve(N, Ps):
    cs = set()
    cs.add(0)
    for p in Ps:
        d = set()
        for c in cs:
            d.add(c + p)
        cs.update(d)
    return len(cs)

if __name__ == "__main__":
    N = int(input())
    Ps = [int(c) for c in input().split(" ")]
    print(solve(N, Ps))

練習のためDPを考える。実装はどうあれ、思考が大事。

まず、順番はどうあれ、すべてを処理しそうなので、0番目からi番目を処理することを考える。

ここで以下のように考える。

「i番目において何通り存在するかを計算しようとしたとき、i-1番目までの処理結果として何が出てくれば、計算できるか。」

今回は、i-1番目までに通り数が出てくるだけではi番目は計算できない。 そこでi-1番目までに作れるすべての数値のSetが必要だと思いつく。 つまりは「i-1番目において、特定の数値が作れるか否か」の情報があればi番目に役に立てられる。

ではi-1番目までの結果をどう得るかで配列と再帰の2パターンで実装する。

配列

def dp_solve1(N, Ps):
    dp = [[False for p in range(sum(Ps) + 1)] for i in range(N + 1)]
    dp[0][0] = True
    for i in range(1, N + 1):
        for j in range(sum(Ps) + 1):
            if j - Ps[i - 1] >= 0:
                dp[i][j] = dp[i - 1][j - Ps[i - 1]] or dp[i - 1][j]
            else:
                dp[i][j] = dp[i - 1][j]
    return sum(dp[N])

配列の場合は漸化式となり、i-1番目の配列に、i-1番目までの結果が入っており、これを元にi番目の配列を作る。 「i番目においてjが作れるか否か」をbooleanとしてdp[i][j]に入れると、

  1. i-1番目においてjが作れていれば、「p[i]を取らない」選択によってi番目でも作れる

  2. i-1番目においてj - p[i]が作れていれば、「p[i]を取る」選択によってi番目でも作れる。

これが大方針となる。初期条件としてはdp[0][0] = True

実装面の注意として

  1. 初期条件は問題を考える配列のさらにひとつ前の配列に存在する必要がある。

  2. 作れる数の最大範囲の分だけ、jの範囲が必要

  3. j - p[i]が0以下のときは作れない。

次に、再帰を考える、

def dp_solve2(N, Ps):
    def dp(i):
        if i == 0:
            return set([0, Ps[i]])
        else:
            prev = dp(i - 1)
            return set([ps + pv for ps in [0, Ps[i]] for pv in prev])
    return len(dp(N - 1))

再帰は「i-1番目において、特定の数値が作れるか否か」をダイレクトに「作れる数値すべて」としてSetとして返す。 一番小さい問題として、0番目だけで考えると、単純に取るか取らないかの2パターンを返し、 i-1番目の結果からさらに、取るか取らないかの2パターン分をすべて計算、Setにして返す。

実装面の注意として

  1. dpの最初のコールはPの問題の添え字なので、Nではなく、N-1。

  2. dpは2つ以上の引数であることが多いが1次元配列でも立派にDP。2つであることを前提に考えない。

実装面では添え字の細かい部分まで気を張る必要があるね。