Kickstart2020 RoundA: Plates

2020 RoundAのPlatesの問題を実装してみました。
実装したものは、こちらにあります。

問題の要約

https://codingcompetitions.withgoogle.com/kickstart/round/000000000019ffc7/00000000001d40bb

N個のスタックに、それぞれK個の皿があり、それぞれの皿には正の値（beauty value）が書かれている。ここからP個の皿を取り出した場合、その和に関して最大値を求めたいという問題です。

アルゴリズム1(TLE)

N個のスタックから皿を取り出す全パターンを考えて、そのパターンのうち、取り出した数がP個である組み合わせを対象に、和の最大値を求めれば問題を解決できそうです。
この時の全パターンは、スタック１からK個を取り出すそれぞれについて、スタック２からK個を取り出す組み合わせがあり、さらにスタック３からK個を取り出す組み合わせがあり、...、最後スタックNまでK個を取り出す場合の組み合わせとして列挙できるので、$K^N$通りの組み合わせがあります。
このアルゴリズムでは、Test set1は$1 \leq N \leq 3$なのでPass出来ますが、$1 \leq N \leq 50$のTest set2ではTLEとなってしまいます。（$1 \leq K \leq 30$）

アルゴリズム２(Passed)

最終的に求めたいのは、N個のスタックからP個の皿を取り出した時の和の最大値なのですが、その途中の状態を考えてみます。$\rm{i}$番目までのスタックから$\rm{j}$個の皿を取り出した時の最大値がdp[i][j]として計算されるとすると、dp[N][P]がN番目までのスタックからP個の皿を取り出した時の最大値で、求めたい解になります。動的計画法の考え方を使って、途中の最大値を計算、更新していけば、組み合わせで求めるよりも大幅に計算する回数が減りそうです。
dp[i][j]の時の最大値は、$0 \leq \rm{x} \leq \rm{min(j, K)}$のそれぞれのxについて、i番目のスタックのx個分の合計と、i-1番目までのスタックでj-x個分の最大値の和として更新します（i番目のスタックからx個取り出し、それ以外のスタックからj-x個取り出した時の最大値の合計で、最大のもの）。

実装

以下のような実装となりました。dp[i][j]の更新の時、累積和（cum_sum）を用いていますが、cum_sumの配列サイズは(N)で、dpの配列サイズは(N+1, P+1)なので、dp[i]に対応するのはcum_sum[i-1]となっています。計算量は$O(N*K*P)$です。

def cumulative_sum(stack, N, K):
    cum_sum = []
    for i in range(N):
        cum_sum_i = [0 for _ in range(K + 1)]
        for j in range(K):
            cum_sum_i[j + 1] = cum_sum_i[j] + stack[i][j]
        cum_sum.append(cum_sum_i)
    return cum_sum
    
def plates(stack, N, P, K):
    dp = [[0] * (P + 1) for _ in range(N + 1)]
    cum_sum = cumulative_sum(stack, N, K)
    for i in range(1, N + 1):
        for j in range(0, P + 1):
            dp[i][j] = 0
            for x in range(min(j, K) + 1):
                dp[i][j] = max(dp[i][j], cum_sum[i - 1][x] + dp[i - 1][j - x])
    return dp[N][P]

T = int(input())
for t in range(1, T + 1):
    N, K, P = list(map(int, input().split()))
    stack = []
    for _ in range(N):
        stack.append(list(map(int, input().split())))
    print('Case #{}: {}'.format(t, plates(stack, N, P, K)))