RGB Coloring [AtCoder Grand Contest 025 B]

https://beta.atcoder.jp/contests/agc025/tasks/agc025_b

前提知識

• modありのaCb

考察

1. （コンビネーションを使うというのは事前に見てしまっていた）
2. まずは塗られる場所というのは特に問題にはならない
3. よって、個数が赤色r、緑色g、青色bとするとrA+g(A+B)+bB=Kとなればいい
4. なにか全探索をするものを探すのだが、これだとr,gの2つを全探索する必要がある
5. 等式を変形するrA+g(A+B)+bB = (r+g)A+(g+b)B=K
6. rgbの列挙からx=r+g、y=g+bの列挙として考えると、xだけ列挙すればyは一意に定まる
7. そこでA点取れるブロックの個数をx, B点取れるブロックの個数をyとすると、組み合わせ数はC(N,x)*C(N,y)となる
8. これはO(1)で計算できるので、ACできそう

解説

https://beta.atcoder.jp/contests/agc025/submissions/2618418

A点取れるブロックの個数をx, B点取れるブロックの個数をyとする。
すると、

• xA+yB=K
• 0≦x,y≦N

を満たすx,yの組で作られる彩色が条件を満たす。
よって、以上の条件を満たすx,yの組で、その条件を満たす彩色の組み合わせ数C(N,x)*C(N,y)を全て足していけばいい。
　
C(N,x)を高速に求めるには、このサイトの方法を参考にしよう。
愚直に計算していてはO(N)かかるので、以上の方法でO(1)で計算しよう。
概要としては、
0. nCk mod p(pは素数)を計算するとする
1. nCk = n! / k! * (n - k)!であるので、n! mod pを事前に計算しておく
2. 問題はk!の除算であるが、(k!)^-1をフェルマーの小定理と繰り返し二乗法で事前計算しておく
3. よって、nCk = n! * (k!)^-1 * (n-k)!は全ての要素が事前計算されているので、O(1)

ll N, A, B, K;
Comb<mint, 1010101> com;
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> A >> B >> K;
 
    mint ans = 0;
    rep(x, 0, N + 1) {
        ll d = K - A * x;
        if (0 <= d and d % B == 0) {
            ll y = d / B;
 
            if (y <= N) {
                ans += com.aCb(N, x) * com.aCb(N, y);
            }
        }
    }
    cout << ans << endl;
}