日々ｄｒｄｒする人のメモ

今日も明日もｄｒｄｒ

トップ > ProCon > AtCoder: Tenka1 Programmer Contest 2019 - D問題: Three Colors

AtCoder: Tenka1 Programmer Contest 2019 - D問題: Three Colors

ProCon AtCoder

コンテスト中に解けなかった。だめだめだった。

atcoder.jp

問題概要

{N}個の整数が与えられ、 {i}番目の数は {a_i}である。

各数について赤、青、緑の3色で塗り分けることを考える。
そして塗り分けた {N}個の数の中で、その色で塗られた整数の総和をそれぞれ {R, G, B}とする。

全ての塗り分け方の中で、辺の長さが {R, G, B}となる三角形が存在する塗り分け方は何通りあるかをmod 998244353で求めよ。

制約

  • {1 \le N \le 300}
  • 各 {i}について {1 \le a_i \le 300}

解法

DPによる計算。

まず、 {R \ge G, B}を仮定する。( {G \ge R, B}の場合と {B \ge R, G}の場合についても同様に計算できる)
こう仮定することで、三角形の存在条件は {R < G + B}を考えればよくなる。

ここで、全ての塗り分け方( {3^N}通り)から存在しない通り数を引くことで求める。
つまり、 {G + B \le R}となる組み合わせを求める。

ここで、以下のdpを考える。

  • {dp[i][s] := } {i}番目までの数で {R + G}の総和が {s}の時の通り数

DPの伝搬は以下のようにする。

  • {dp[i+1][s] \leftarrow dp[i][s]}
  • {dp[i+1][s+a_i] \leftarrow 2 \times dp[i][s]}

そして、 {\displaystyle S = \sum_{i=1}^N a_i}とした時、
{i \le S - i} を満たす {dp[N][i]} の総和を解の値から引けばよい。


同様に {G \ge R, B} と {B \ge R, G} の場合についても解の値から引けばよいが、
その結果、

  • {R = G = S/2, B = 0}
  • {R = B = S/2, G = 0}
  • {G = B = S/2, R = 0}

の場合がそれぞれ1回多く引かれているため、それらを足し直す必要がある。

足し直す通り数は、 {N}個の整数の総和をちょうど半分ずつに分ける通り数として、別のDPで求めればよい。

計算量は {\displaystyle O(N^2 \max_i a_i)}

実装

提出コード (PyPy3): Submission #5067923 - Tenka1 Programmer Contest 2019

N, *A = map(int, open(0).read().split())
MOD = 998244353

S = sum(A)
# 三角形が存在しない通り数を求めるためのDP
dp0 = [0]*(S+1)
dp0[0] = 1
for a in A:
    for i in range(S, a-1, -1):
        dp0[i] = (dp0[i] + dp0[i-a]*2) % MOD

# 足し直す分のDP
dp1 = [0]*(S+1)
dp1[0] = 1
for a in A:
    for i in range(S, a-1, -1):
        dp1[i] = (dp1[i] + dp1[i-a]) % MOD

ans = pow(3, N, MOD)
# 解から三角形が存在しない通り数を引く
for L in range(0, S//2+1):
    ans = (ans - dp0[L]*3) % MOD
# 引きすぎた分を足す
if S % 2 == 0:
    ans = (ans + dp1[S//2]*3) % MOD

print(ans)