😎

Kickstart2020 RoundA: Plates

2 min read

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

問題の要約

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

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

アルゴリズム1(TLE)

N個のスタックから皿を取り出す全パターンを考えて、そのパターンのうち、取り出した数がP個である組み合わせを対象に、和の最大値を求めれば問題を解決できそうです。
この時の全パターンは、スタック1からK個を取り出すそれぞれについて、スタック2からK個を取り出す組み合わせがあり、さらにスタック3から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

アルゴリズム2(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)))

Discussion

ログインするとコメントできます