Problem p03332 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

using namespace std;

typedef long long int ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int INF = 1000000000;
const ll LINF = (ll)4000000000000000000;
const ll MOD = (ll)998244353;
const double PI = acos(1.0);
const int limit = 100010;

#define REP(i, m, n) for (ll i = m; i < (ll)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define MP make_pair
#define YES(n) cout << ((n) ? "YES" : "NO") << endl
#define Yes(n) cout << ((n) ? "Yes" : "No") << endl
#define Possible(n) cout << ((n) ? "Possible" : "Impossible") << endl
#define NP(v) next_permutation(v.begin(), v.end())
#define smaller_queue priority_queue<ll, vector<ll>, greater<ll>>

//------------------------------------------------------

const ll MAX_N = 3 * limit;
ll fact[MAX_N];
ll fact_inv[MAX_N];

ll bin_pow_mod(ll x, ll y) { // x^y mod
  if (x == 0)
    return 0;
  ll prod = 1;
  while (y > 0) {
    if (y % 2 == 1)
      prod = (prod * x) % MOD;
    x = (x * x) % MOD;
    y /= 2;
  }
  return prod % MOD;
}

void init_fact() {
  fact[0] = fact_inv[0] = 1;

  REP(i, 1, MAX_N) {
    fact[i] = (fact[i - 1] * i) % MOD;
    fact_inv[i] = bin_pow_mod(fact[i], MOD - 2);
  }
}

ll comb_mod(ll n, ll r) {
  if (n < r)
    return 0;
  else
    return ((fact[n] * fact_inv[n - r]) % MOD * fact_inv[r]) % MOD;
}

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

  init_fact();

  vector<pll> v;

  REP(m, 0, n + 1) {
    ll bi_n = k - a * m;
    if (bi_n >= 0 && bi_n % b == 0 && bi_n / b <= n)
      v.push_back(MP(m, bi_n / b));
  }

  ll ans = 0;

  for (auto p : v) {
    ll ai = p.first;
    ll bi = p.second;

    ll green = 0;

    REP(g, 0, min(ai, bi) + 1) {
      ll r = ai - g;
      ll b = bi - g;
      ans += comb_mod(n, r) * comb_mod(n - r, b) % MOD *
             comb_mod(n - r - b, g) % MOD;
      ans %= MOD;
    }
  }

  cout << ans << endl;

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>

using namespace std;

typedef long long int ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int INF = 1000000000;
const ll LINF = (ll)4000000000000000000;
const ll MOD = (ll)998244353;
const double PI = acos(1.0);
const int limit = 100010;

#define REP(i, m, n) for (ll i = m; i < (ll)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define MP make_pair
#define YES(n) cout << ((n) ? "YES" : "NO") << endl
#define Yes(n) cout << ((n) ? "Yes" : "No") << endl
#define Possible(n) cout << ((n) ? "Possible" : "Impossible") << endl
#define NP(v) next_permutation(v.begin(), v.end())
#define smaller_queue priority_queue<ll, vector<ll>, greater<ll>>

//------------------------------------------------------

const ll MAX_N = 3 * limit;
ll fact[MAX_N];
ll fact_inv[MAX_N];

ll bin_pow_mod(ll x, ll y) { // x^y mod
  if (x == 0)
    return 0;
  ll prod = 1;
  while (y > 0) {
    if (y % 2 == 1)
      prod = (prod * x) % MOD;
    x = (x * x) % MOD;
    y /= 2;
  }
  return prod % MOD;
}

void init_fact() {
  fact[0] = fact_inv[0] = 1;

  REP(i, 1, MAX_N) {
    fact[i] = (fact[i - 1] * i) % MOD;
    fact_inv[i] = bin_pow_mod(fact[i], MOD - 2);
  }
}

ll comb_mod(ll n, ll r) {
  if (n < r)
    return 0;
  else
    return ((fact[n] * fact_inv[n - r]) % MOD * fact_inv[r]) % MOD;
}

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

  init_fact();

  vector<pll> v;

  REP(m, 0, n + 1) {
    ll bi_n = k - a * m;
    if (bi_n >= 0 && bi_n % b == 0 && bi_n / b <= n)
      v.push_back(MP(m, bi_n / b));
  }

  ll ans = 0;

  for (auto p : v) {
    ll ai = p.first;
    ll bi = p.second;

    ans += comb_mod(n, ai) * comb_mod(n, bi) % MOD;
    ans %= MOD;
  }

  cout << ans << endl;

  return 0;
}