Tenka1 Programmer Contest 2019 D - Three Colors

問題ページ

考えたこと

• 部分和問題っぽくて貪欲は無理そうだし制約小さいしまあDP
• dp[i番目まで][R=j][G=j] みたいなDPを考える
• R,G,Bは90000くらいまで取りうるのでどう考えても無理
• 三角形の成立条件を考える
• R+G>B && R+B>G && B+G>R
• どれか一つ和を持っておけば済むとかがあると嬉しいけど思いつかない
• 条件の否定を取って補集合の方を数えることにしてみる
• R+G<=B || R+B<=G || B+G<=R になる
• 和集合を数えたいので包除原理っぽい
• R+G<=Bだけを満たすのを数えるだけならDPすればできる
• 対称性があるので3倍すればよさそう
• 2つ条件を満たす場合を数え上げたい
• R+G<=B && R+B<=G
• いろいろ考えてしばらく迷走する
• これ満たすのB=G, R=0 しかない
• R=0 or G=0 or B=0 だったら三角形ができることはない
• R>0 && G>0 && B>0 && R+G<=B のパターンを数え上げる
• 対称性でこれを3倍する
• 全体(R!=0&&G!=0&&B!=0の塗り方)からこの値を引けば答えになる
• dp[i番目まで][j=R+G][k=(R>0)][l=(G>0)] = (組み合わせ数) としてDP

modintバグってたのとmake_vをグローバルにしたら爆速になったのとかが虚無だった

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
// #define int ll
using PII = pair<ll, ll>;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) x.begin(), x.end()

template<typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template<typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template<typename T> bool IN(T a, T b, T x) { return a<=x&&x<b; }
template<typename T> T ceil(T a, T b) { return a/b + !!(a%b); }

template<typename T> vector<T> make_v(size_t a) { return vector<T>(a); }
template<typename T,typename... Ts>
auto make_v(size_t a,Ts... ts) {
    return vector<decltype(make_v<T>(ts...))>(a,make_v<T>(ts...));
}
template<typename T,typename V> typename enable_if<is_class<T>::value==0>::type
fill_v(T &t, const V &v) { t=v; }
template<typename T,typename V> typename enable_if<is_class<T>::value!=0>::type
fill_v(T &t, const V &v ) { for(auto &e:t) fill_v(e,v); }

template<class S,class T>
ostream &operator <<(ostream& out,const pair<S,T>& a){
    out<<'('<<a.first<<','<<a.second<<')'; return out;
}
template<class T>
ostream &operator <<(ostream& out,const vector<T>& a){
    out<<'['; for(T i: a) {out<<i<<',';} out<<']'; return out;
}

int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0}; // DRUL
const int INF = 1<<30;
const ll LLINF = 1LL<<60;
// const ll MOD = 1000000007;

template<ll MOD>
struct modint {
    ll x;
    modint(): x(0) {}
    modint(ll y) : x(y>=0 ? y%MOD : y%MOD+MOD) {}
    static constexpr ll mod() { return MOD; }
    // e乗
    modint pow(ll e) {
        ll a = 1, p = x;
        while(e > 0) {
            if(e%2 == 0) {p = (p*p) % MOD; e /= 2;}
            else {a = (a*p) % MOD; e--;}
        }
        return modint(a);
    }
    modint inv() const {
        ll a=x, b=MOD, u=1, y=1, v=0, z=0;
        while(a) {
            ll q = b/a;
            swap(z -= q*u, u);
            swap(y -= q*v, v);
            swap(b -= q*a, a);
        }
        return z;
    }
    // Comparators
    bool operator <(modint b) { return x < b.x; }
    bool operator >(modint b) { return x > b.x; }
    bool operator<=(modint b) { return x <= b.x; }
    bool operator>=(modint b) { return x >= b.x; }
    bool operator!=(modint b) { return x != b.x; }
    bool operator==(modint b) { return x == b.x; }
    // Basic Operations
    modint operator+(modint r) const { return modint(*this) += r; }
    modint operator-(modint r) const { return modint(*this) -= r; }
    modint operator*(modint r) const { return modint(*this) *= r; }
    modint operator/(modint r) const { return modint(*this) /= r; }
    modint &operator+=(modint r) {
        if((x += r.x) >= MOD) x -= MOD;
        return *this;
    }
    modint &operator-=(modint r) {
        if((x -= r.x) < 0) x += MOD;
        return *this;
    }
    modint &operator*=(modint r) { x = x * r.x % MOD; return *this; }
    modint &operator/=(modint r) { return *this *= r.inv(); }
    // increment, decrement
    modint operator++() { x++; return *this; }
    modint operator++(signed) { modint t = *this; x++; return t; }
    modint operator--() { x--; return *this; }
    modint operator--(signed) { modint t = *this; x--; return t; }
};
using mint = modint<998244353>;
template<class T> mint operator*(T l, mint r) { return mint(l) *= r; }
template<class T> mint operator+(T l, mint r) { return mint(l) += r; }
template<class T> mint operator-(T l, mint r) { return mint(l) -= r; }
template<class T> mint operator/(T l, mint r) { return mint(l) /= r; }
template<class T> bool operator==(T l, mint r) { return mint(l) == r; }
template<class T> bool operator!=(T l, mint r) { return mint(l) != r; }
// Input/Output
ostream &operator<<(ostream& os, mint a) { return os << a.x; }
istream &operator>>(istream& is, mint &a) { return is >> a.x; }
string to_frac(mint v) {
    static map<ll, PII> mp;
    if(mp.empty()) {
        mp[0] = mp[mint::mod()] = {0, 1};
        FOR(i, 2, 1001) FOR(j, 1, i) if(__gcd(i, j) == 1) {
            mp[(mint(i) / j).x] = {i, j};
        }
    }
    auto itr = mp.lower_bound(v.x);
    if(itr != mp.begin() && v.x - prev(itr)->first < itr->first - v.x) --itr;
    string ret = to_string(itr->second.first + itr->second.second * ((int)v.x - itr->first));
    if(itr->second.second > 1) {
        ret += '/';
        ret += to_string(itr->second.second);
    }
    return ret;
}

mint dp[300][90001][2][2];
signed main(void)
{
    cin.tie(0);
    ios::sync_with_stdio(false);

    ll n;
    cin >> n;
    ll sum = 0;
    vector<ll> a(n);
    REP(i, n) cin >> a[i], sum += a[i];

    dp[0][0][0][0] = 1;
    dp[0][a[0]][1][0] = 1;
    dp[0][a[0]][0][1] = 1;
    FOR(i, 1, n) REP(j, sum+1) REP(k, 2) REP(l, 2) {
        if(j>=a[i]) {
            dp[i][j][k|1][l] += dp[i-1][j-a[i]][k][l];
            dp[i][j][k][l|1] += dp[i-1][j-a[i]][k][l];
        }
        dp[i][j][k][l] += dp[i-1][j][k][l];
    }

    mint ret = 0;
    FOR(i, 1, sum) if(i <= sum-i) ret += dp[n-1][i][1][1];
    cout << mint(3).pow(n) - mint(2).pow(n)*3 + 3 - ret*3 << endl;

    return 0;
}