AtCoder Grand Contest 025 B - RGB Coloring

問題

　最初全て無色である個のブロックを赤緑青の3色で塗り分けることを考える。各ブロックがそれぞれ赤であれば点、緑であれば点、青であれば点、無色のままなら点として、個のブロックの合計得点がになる塗り方は何通りか求めよ。

解法

　解説PDFの通り。

　緑で塗られたブロックを赤と青で同時に塗られたと考えれば赤と青を独立にそれぞれ個、個塗る問題に帰着する。 $$ K = i A + j B $$

としてを全探索すればと求まるのでで解ける。

反省

　さすがにたった2か月前のコンテストということで解法を覚えてしまっていた。10分ほどでAC。コンテスト中は解けなかったわけだけど、なかなかこういう発想の転換ができてこない。緑の点数がということで露骨に怪しいとは思っていた記憶があるが、こうやって扱うと簡単に解けるんだなぁ。

コード

#include"bits/stdc++.h"
using namespace std;
using ll = long long;

constexpr ll MOD = 998244353;

vector<ll> fact, inv_fact;

ll combination(ll n, ll m) {
    return fact[n] * inv_fact[n - m] % MOD * inv_fact[m] % MOD;
}

ll POW(ll n, ll m) {
    ll result = 1;
    while (m) {
        if (m & 1) {
            result *= n;
            result %= MOD;
        }

        m /= 2;
        n *= n;
        n %= MOD;
    }
    return result;
}

int main() {
    ll N, A, B, K;
    cin >> N >> A >> B >> K;

    //緑→紫として紫は赤と青が同時に塗られている場所とすると
    //N個のブロックから赤をi個,青をj個それぞれ独立に選んで色
    //を塗ることとして良い

    //fact, inv_factの準備
    fact.resize(N + 1, 1);
    inv_fact.resize(N + 1, 1);
    for (ll i = 2; i <= N; i++) {
        fact[i] = i * fact[i - 1] % MOD;
        inv_fact[i] = POW(fact[i], MOD - 2);
        assert(fact[i] * inv_fact[i] % MOD == 1);
    }

    ll ans = 0;
    for (ll i = 0; i <= N; i++) {
        if (i * A > K) {
            continue;
        }

        if ((K - i * A) % B == 0) {
            ll j = (K - i * A) / B;
            if (j > N) {
                continue;
            }
            ans += combination(N, i) * combination(N, j) % MOD;
            ans %= MOD;
        }
    }

    cout << ans << endl;
}