Problem p03332 (C++)

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

#define INF 1000010000
#define nl '\n'
#define pb push_back
#define ppb pop_back
#define mp make_pair
#define fi first
#define se second
#define pii pair<int, int>
#define pdd pair<double, double>
#define all(c) (c).begin(), (c).end()
#define SORT(c) sort(all(c))
#define sz(c) (c).size()
#define rep(i, n) for (int i = 0; i < n; ++i)
#define repi(i, n) for (int i = 1; i <= n; ++i)
#define repn(i, n) for (int i = n - 1; i >= 0; --i)
#define repf(j, i, n) for (int j = i; j < n; ++j)
#define die(s)                                                                 \
  { std::cout << s << nl; }
#define dier(s)                                                                \
  {                                                                            \
    std::cout << s;                                                            \
    return 0;                                                                  \
  }
#define dbg(var)                                                               \
  { std::cout << #var << "  = " << var << nl; }
#define vi vector<int>
typedef long long ll;
#define MOD 998244353LL
using namespace std;

struct CNK {
  const ll mod = 998244353LL;
  int mx;

private:
  vector<ll> fac, ifac;

public:
  CNK(int max) {
    mx = max;

    fac.resize(max + 1);
    ifac.resize(max + 1);

    fac[0] = ifac[0] = 1;

    for (int i = 1; i <= max; ++i) {
      fac[i] = 1LL * fac[i - 1] * i % mod;
      ifac[i] = fpow(fac[i], mod - 2);
    }
  }

  ll fpow(ll a, ll b) {
    ll ans = 1;
    while (b > 0) {
      if (b & 1) {
        ans = (1LL * ans * a) % mod;
      }
      a = (1LL * a * a) % mod;
      b /= 2;
    }
    return ans;
  }

  ll cnk(int n, int k) {
    assert(mx >= n);
    if (n < 0 || k < 0 || k > n)
      return 0;
    return 1LL * fac[n] * ifac[n - k] % mod * ifac[k] % mod;
  }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.precision(0);

  int n, a, b;
  ll k;
  cin >> n >> a >> b >> k;
  ll ans = 0;
  CNK cnk((int)3e5 + 16);
  // cout << "green red blue\n";
  rep(x1, n + 1) {
    int x2 = (k - x1 * a) / b;
    if (1LL * a * x1 + 1LL * b * x2 != k)
      continue;
    //	cout << x1 << " " << x2 << endl;
    if (x1 + x2 > n) {
      /*	int gr = x1 + x2 - n;
              ll cur = cnk.cnk(n , gr);
              cur *= cnk.cnk(n - gr , x1 - gr);
              cur %= MOD;
              cur *= cnk.cnk(n - gr - (x1 - gr) , x2 - gr);
              cur %= MOD;
              ans = (ans + cur) % MOD;*/
      rep(i, 1 + min(x1, x2)) {
        //	cout << i << " " << x1 - i << " " << x2 - i << endl;
        ll cur = cnk.cnk(n, i);
        cur *= cnk.cnk(n - i, x1 - i);
        cur %= MOD;
        cur *= cnk.cnk(n - i - (x1 - i), x2 - i);
        cur %= MOD;
        ans = (ans + cur) % MOD;
      }
    } else {
      rep(i, 1 + min(x1, x2)) {
        //	cout << i << " " << x1 - i << " " << x2 - i << endl;
        ll cur = cnk.cnk(n, i);
        cur *= cnk.cnk(n - i, x1 - i);
        cur %= MOD;
        cur *= cnk.cnk(n - i - (x1 - i), x2 - i);
        cur %= MOD;
        ans = (ans + cur) % MOD;
      }
    }
  }
  cout << ans % MOD << endl;

  return 0;
}

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

#define INF 1000010000
#define nl '\n'
#define pb push_back
#define ppb pop_back
#define mp make_pair
#define fi first
#define se second
#define pii pair<int, int>
#define pdd pair<double, double>
#define all(c) (c).begin(), (c).end()
#define SORT(c) sort(all(c))
#define sz(c) (c).size()
#define rep(i, n) for (int i = 0; i < n; ++i)
#define repi(i, n) for (int i = 1; i <= n; ++i)
#define repn(i, n) for (int i = n - 1; i >= 0; --i)
#define repf(j, i, n) for (int j = i; j < n; ++j)
#define die(s)                                                                 \
  { std::cout << s << nl; }
#define dier(s)                                                                \
  {                                                                            \
    std::cout << s;                                                            \
    return 0;                                                                  \
  }
#define dbg(var)                                                               \
  { std::cout << #var << "  = " << var << nl; }
#define vi vector<int>
typedef long long ll;
#define MOD 998244353LL
using namespace std;

struct CNK {
  const ll mod = 998244353LL;
  int mx;

private:
  vector<ll> fac, ifac;

public:
  CNK(int max) {
    mx = max;

    fac.resize(max + 1);
    ifac.resize(max + 1);

    fac[0] = ifac[0] = 1;

    for (int i = 1; i <= max; ++i) {
      fac[i] = 1LL * fac[i - 1] * i % mod;
      ifac[i] = fpow(fac[i], mod - 2);
    }
  }

  ll fpow(ll a, ll b) {
    ll ans = 1;
    while (b > 0) {
      if (b & 1) {
        ans = (1LL * ans * a) % mod;
      }
      a = (1LL * a * a) % mod;
      b /= 2;
    }
    return ans;
  }

  ll cnk(int n, int k) {
    assert(mx >= n);
    if (n < 0 || k < 0 || k > n)
      return 0;
    return 1LL * fac[n] * ifac[n - k] % mod * ifac[k] % mod;
  }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.precision(0);

  int n, a, b;
  ll k;
  cin >> n >> a >> b >> k;
  ll ans = 0;
  CNK cnk((int)3e5 + 16);
  // cout << "green red blue\n";
  rep(x1, n + 1) {
    int x2 = (k - x1 * a) / b;
    if (1LL * a * x1 + 1LL * b * x2 != k)
      continue;
    ans += (cnk.cnk(n, x1) * cnk.cnk(n, x2)) % MOD;
    ans %= MOD;
  }
  cout << ans % MOD << endl;

  return 0;
}