Tenka1 Programmer Contest 2019 D Three Colors

Tenka1 Programmer Contest 2019 に参加した。

C しか解けなかった。
D の Three Colors が解けそうで解けなかったので、その解説。
問題はこれ
D - Three Colors

解説

ほぼほぼ公式解説の通りだけどね。
 S = \sum_{i=1}^N a_i とする。
三角形の最長の辺が S/2 以上だと、正の面積を持つ三角形にならない。
そうでないものが求めるもの。
R が最長辺と仮定して良い。
そうすると、その 3 倍が求めたいものです。
ここで、2 辺が S/2 に等しい場合をダブルカウントしてしまっていることに注意。
細かいことはプログラムの中に書く。
ちなみに、pypy じゃないと通らなかった。

解答

Q = 998244353
N = int( input())
A = [ int( input()) for _ in range(N)]
S = sum(A)
H = (S+1)//2
dp = [0]*(S+1)
dp[0] = 1
for a in A:
    for s in range(S, a-1, -1):
        dp[s] = (dp[s]*2 + dp[s-a])%Q #dp[s]*2 はR以外においた場合の数。GとBに置く2通りがある。
    for s in range(a-1,-1,-1):
        dp[s] = dp[s]*2%Q
ans = (pow(3, N, Q) - sum(dp[H:])%Q*3%Q)%Q
if S%2 == 0:#R=G=S/2となる場合を数える。
    dp = [0]*(S+1)
    dp[0] = 1
    for a in A:
        for s in range(S,a-1,-1):
            dp[s] = (dp[s] + dp[s-a])%Q
    ans += dp[H]*3 #R=G=S/2, R=B=S/2, B=G=S/2の3通りがあるので
    ans %= Q
print(ans)