公式解説: https://img.atcoder.jp/tenka1-2019/editorial.pdf

問題概要

N 個の整数 $ a_1, a_2, \cdots a_N $ が与えられる．これらの整数を赤，緑，青で塗り分ける．

赤，緑，青で塗られた整数の和を R, G, B とするとき， 3辺の長さが R, G, B であるような正の面積の三角形が存在するような塗り分け方の数を 998244353 で割ったあまりを求めよ．

解法概要

$ \sum_{i=1}^{N} a_i = S $ とする．

正の面積の三角形が存在する条件は，$ R+G > B $ ，$ G+B > R $ ，$ B+R > G $ の全てが成り立つことである．すなわち，一番長い辺の長さが$ \frac{S}{2} $ 以上の場合，正の面積の三角形は存在しない．一番長い辺の長さが$ \frac{S}{2} $ 未満の場合には，正の面積の三角形が存在する (一番長い辺の長さが$ \frac{S}{2} $ 未満ということは他の辺の長さもそれ以下にしなければならず，いずれの辺の長さも 0 にはならないため) ．

ここで $ R \geq G, B $ の場合だけを考えてみる．

$ R \geq \frac{S}{2}$ のとき (すなわち，正の面積の三角形が存在しないとき) $ B,G \leq \frac{S}{2}$ なので， $ R \geq B,G $ が必ず成り立っている．つまり，

$ \mathit{dp}_{i, r} = a_i $まで見て，赤に塗った棒の長さの合計が r の場合の塗り方の通り数

$ \mathit{dp}_{i+1, r} = \mathit{dp}_{i+1, r} + \mathit{dp}_{i,r} * 2 $ (今見ている整数を青または緑に塗る場合に対応)

$ \mathit{dp}_{i+1, r+a_i} = \mathit{dp}_{i+1, r+a_i} + \mathit{dp}_{i,r} $ (今見ている整数を赤に塗る場合に対応)

というような DP をして最後に $ \sum_{r=\frac{S}{2}}^{S} \mathit{dp}_{n,r} $ を求めれば，R が最も長くかつ正の面積の三角形が存在しない塗り方の通り数が求められる．

あとは，緑が最も長い場合と青が最も長い場合が存在するので，解は全ての塗り分け方の通り数から $ 3 \sum_{r=\frac{S}{2}}^{S} \mathit{dp}_{n,r} $ を引いたもの……といきたいところであるが，ここは注意が必要．「赤が最も長い」「緑が最も長い」「青が最も長い」は排他的な条件ではなく，例えば，「赤が最も長い」かつ「緑が最も長い」というような状態が存在する．S が偶数のとき，単純に 3 倍してしまうと， $ R = G = \frac{S}{2}, B = 0 $ というような 2色だけで色分けを行なっていて，かつ和が等しいような場合を 2 回引いてしまうことになる．そのため， S が偶数の場合は 2 色に色分けをしてそれぞれの長さが $ \frac{S}{2} $ となるような塗り分け方の数を数えておかなければならない．これも，例えば赤と緑だけを使うと仮定すると

$ \mathit{dp2}_{i, r} = a_i $まで見て，赤に塗った棒の長さの合計が r の場合の塗り方の通り数

$ \mathit{dp2}_{i+1, r} = \mathit{dp2}_{i+1, r} + \mathit{dp2}_{i,r}$ (今見ている整数を緑に塗る場合に対応)

$ \mathit{dp2}_{i+1, r+a_i} = \mathit{dp2}_{i+1, r+a_i} + \mathit{dp2}_{i,r} $ (今見ている整数を赤に塗る場合に対応)

という，先ほどと似た DP で求められる．

全ての塗り分け方は $ 3^N $ 通りなので，求めるべき解は

$ 3^N - 3 \sum_{r=\frac{S}{2}}^{S} \mathit{dp}_{n,r} $ (N が奇数のとき)

$ 3^N - (3 \sum_{r=\frac{S}{2}}^{S} \mathit{dp}_{n,r} - 3\mathit{dp2}_{n, \frac{r}{2}} ) $ (N が偶数のとき)

となる．

ソースコード

#include <iostream>
#include <vector>
#include <numeric>

using namespace std;

typedef long long ll;

const ll MOD = 998244353;
const int N_MAX = 300;
const int A_MAX = 300;

vector<vector<ll>> dp(N_MAX+1, vector<ll>(N_MAX*A_MAX+1, 0)), dp2(N_MAX+1, vector<ll>(N_MAX*A_MAX+1, 0));

ll modpow(ll a, ll t) {
    ll ret = 1LL;
    while(t){
        if(t & 1LL){
            ret *= a;
            ret %= MOD;
        }
        a *= a;
        a %= MOD;
        t >>= 1;
    }
    return ret;
}

int main() {
    int n;
    cin >> n;
    vector<ll> a(n);
    for(int i=0;i<n;++i){
        cin >> a[i];
    }
    ll s = accumulate(a.begin(), a.end(), 0LL);

    dp[0][0] = 1;
    dp2[0][0] = 1;
    for(int i=0;i<n;++i) {
        for(int r=0;r<=s;++r) {
            if(!dp[i][r]) continue;
            dp[i+1][r] += 2 * dp[i][r];
            dp[i+1][r] %= MOD;
            dp[i+1][r+a[i]] += dp[i][r];
            dp[i+1][r] %= MOD;
            dp2[i+1][r+a[i]] += dp2[i][r];
            dp2[i+1][r+a[i]] %= MOD;
            dp2[i+1][r] += dp2[i][r];
            dp2[i+1][r] %= MOD;
        }
    }

    ll ans = 0LL;
    for(int r=(s+1)/2;r<=s;++r) {
        ans += dp[n][r];
        ans %= MOD;
    }
    if(s % 2 == 0) ans += MOD - dp2[n][s/2];
    ans = (3 * ans) % MOD;
    ans = (modpow(3LL, n) + MOD - ans) % MOD;

    cout << ans << endl;
}

所感

余事象を数えるぞ！ってところまでしか本番は思いつかなかった．ずっと「赤に塗った整数の集合の大きさがわからないと……」とか寝ぼけたこと言ってましたね．

(2019/04/24 追記) 肝心の「求めるべき解は〜」が余事象の個数のままになっていたので直しました

NoiminのNoise

Tenka1 Programmer Contest 2019 D - Three Colors

問題概要

解法概要

ソースコード

所感