Tenka1 Programmer Contest 2019 D問題：Three Colors

問題. Three Colors

$N$ 個の整数 $s_1, \ldots, s_N$ が与えられる．与えられたすべての整数を赤，緑，青の3色のいずれかで塗るとき，次の条件を満たすような塗り方が何通りあるかを答えよ．

赤，緑，青で塗られた整数の和をそれぞれ R, G, B としたとき，3辺の長さを R, G, B とした三角形が存在する

制約： $3 \le N \le 300$ ，　 $1 \le s_i \le 300$

解法．包除原理＋動的計画法

自力で解けなかったので Editorial を参考に解いた．

3色で塗り分けたときの各色の整数の総和を $a, b, c$ とする．このとき，3辺の長さが $a, b, c$ となる三角形が存在するための必要十分条件は，
　 $a + b > c$ かつ $b + c > a$ かつ $c + a > b$
である（三角形の成立条件とその証明 | 高校数学の美しい物語）．この条件と同値な条件として，
　 $\frac{S}{2}> c$ かつ $\frac{S}{2} > a$ かつ $\frac{S}{2} > b$
がある．ただし， $S = \sum_{i = 1}^N s_i$ とする（各不等式にそれぞれ $c, a, b$ を足すと成り立つ）．ここで，

$U :=$ すべての塗り分け方
$R :=$ 赤色で塗られた整数の総和が $S / 2$ 以上となる塗り分け方
$G :=$ 緑色で塗られた整数の総和が $S / 2$ 以上となる塗り分け方
$B :=$ 青色で塗られた整数の総和が $S / 2$ 以上となる塗り分け方

とおくと，答えは包除原理より，
　　 $|U| - |R| - |G| - |B| + |R \cap G| + |R \cap B| + |G \cap B| - |R \cap G \cap B|$
となる． $R, G, B$ を $\frac{S}{2} > a$ の補集合として定義したのは共通部分の計算がしやすいからである．
定義から， $|R| = |G| = |B|$ かつ $|R \cap G| = |R \cap B| = |G \cap B|$ となるので，
　　 $|U| - 3 |R| + 2 |R \cap G| - |R \cap G \cap B|$
を求める． $|U|$ は各整数を赤，緑または青への塗り方の総数なので $3^N$ 通りであり， $|R \cap G \cap B|$ はそのような塗り方が存在しないので 0 通りである． $|R|$ と $|R \cap G|$ はそれぞれ異なる動的計画法で求めることができる．

計算時間： $O(N^2 \max(s_i))$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
constexpr ll MOD = 998244353;


ll Solve() {
    int n;
    cin >> n;

    vector<int> a(n);
    for (auto &a_i : a) cin >> a_i;

    const int SUM = accumulate(a.begin(), a.end(), 0);
    vector<vector<ll>> dp1(n + 1, vector<ll>(SUM + 1));
    vector<vector<ll>> dp2(n + 1, vector<ll>(SUM + 1));
    dp1[0][0] = dp2[0][0] = 1;
    for (int i = 0, sum = 0; i < n; ++i) {
        for (int j = 0; j <= sum; ++j) {
            if (dp1[i][j] != 0) {
                (dp1[i + 1][j] += 2 * dp1[i][j]) %= MOD;
                (dp1[i + 1][j + a[i]] += dp1[i][j]) %= MOD;

                (dp2[i + 1][j] += dp2[i][j]) %= MOD;
                (dp2[i + 1][j + a[i]] += dp2[i][j]) %= MOD;
            }
        }
        sum += a[i];
    }

    ll res = 1;
    for (int i = 0; i < n; ++i) (res *= 3) %= MOD;

    ll cnt1 = 0;
    for (int i = (SUM + 1) / 2; i <= SUM; ++i)
        (cnt1 += dp1[n][i]) %= MOD;
    (cnt1 *= 3) %= MOD;
    (cnt1 = MOD - cnt1) %= MOD;

    ll cnt2 = 0;
    if (SUM % 2 == 0) {
        (cnt2 = 3 * dp2[n][SUM / 2]) %= MOD;
    }

    return (res + cnt1 + cnt2) % MOD;
}

int main() {
    cin.tie(0); ios::sync_with_stdio(false);

    cout << Solve() << endl;

    return 0;
}