AtCoder Grand Contest 025 B - RGB Coloring

• 問題

B - RGB Coloring

• 解法

緑色で塗られたマスは点もらえることから、緑色のマス目は赤色と青色の両方が塗られたマス目だと思うことができる。赤色、青色、緑色に塗られたマス目の個数をそれぞれとすると、このときの得点は点である。ここで であることに注意する。得点が点の塗り方の個数を数えればよいので、結局この問題は を求めれば良いことになる。に対しは一意に決まるのでに関して全探索すればこの問題はの計算量時間で解ける。は事前に計算しておきます。

大きいに対してを求める方法は蟻本を見るかググってみてください。

#include <iostream>
#include <string>
#include <algorithm>
#include <cstdio>
#include <vector>
#include <queue>
#include <set>
#include <map>
#include <numeric>
#include <cmath>

using namespace std;

typedef long long int ll;

#define all(x) x.begin(),x.end()

const ll mod = 998244353;
const ll INF = 1e9;
const ll MAXN = 1e9;

vector<ll> fact,inv;

ll f(ll a,ll b, ll p)  // a^b mod pを求める
{
  if (b == 0)
  {
    return 1;
  }else if (b % 2 == 0)
  {
    ll d = f(a,(ll)b/2,p);  //a^b = a^(2k) = (a^k)^2  k = b / 2
    return d * d % mod;
  }else
  {
    return a * f(a,b-1,p) % mod;
  }
}

void init(ll n)
{
  fact.resize(n+1,(ll)1);
  inv.resize(n+1,(ll)1);

  for(int i = 2; i <= n; i++)  //　nまでの階乗を求める
  {
    fact[i] = fact[i-1] * i % mod;
  }
  for(int i = 2; i <= n; i++)  //a^(-1) mod p = a^(p-2) mod p
  {
    inv[i] = f(fact[i],mod-2,mod) % mod;
  }
}

ll comb(ll n,ll k)  //n_C_kを返す
{
  if (n < k)
  {
    return 0;
  }
  return (fact[n] * inv[k] % mod) * inv[n-k] % mod;
}

int main()
{
  ll n,a,b,k;
  cin >> n >> a >> b >> k;

  init(n);
  ll ans = 0;
  for(ll x = 0; x <= n; x++){
    if((k-x*a) % b != 0) continue; //Ax+By=Kを満たす整数yがない
    ll y = (k-x*a)/b;
    if(y < 0 || n < y) continue; //0 <= y <= n
    ans += (comb(n,x) % mod) * comb(n,y) % mod;
    ans %= mod;
  }

  cout << ans << endl;

  return 0;