コンテスト中に解けなかった。だめだめだった。
問題概要
個の整数が与えられ、番目の数はである。
各数について赤、青、緑の3色で塗り分けることを考える。
そして塗り分けた個の数の中で、その色で塗られた整数の総和をそれぞれとする。
全ての塗り分け方の中で、辺の長さがとなる三角形が存在する塗り分け方は何通りあるかをmod 998244353で求めよ。
制約
- 各について
解法
DPによる計算。
まず、を仮定する。(の場合との場合についても同様に計算できる)
こう仮定することで、三角形の存在条件はを考えればよくなる。
ここで、全ての塗り分け方(通り)から存在しない通り数を引くことで求める。
つまり、となる組み合わせを求める。
ここで、以下のdpを考える。
- 番目までの数での総和がの時の通り数
DPの伝搬は以下のようにする。
そして、とした時、
を満たす の総和を解の値から引けばよい。
同様に と の場合についても解の値から引けばよいが、
その結果、
の場合がそれぞれ1回多く引かれているため、それらを足し直す必要がある。
足し直す通り数は、個の整数の総和をちょうど半分ずつに分ける通り数として、別のDPで求めればよい。
計算量は
実装
提出コード (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)