Problem p02662 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, n) for (int i = (a); i < (int)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define FOR(it, c)                                                             \\
  for (__typeof((c).begin()) it = (c).begin(); it != (c).end(); ++it)
#define ALLOF(c) (c).begin(), (c).end()
typedef long long ll;
typedef unsigned long long ull;

template <ll mod, ll primitive_root> class NTT {
  void bit_reverse(vector<ll> &a) {
    int N = a.size();
    int i = 0;
    for (int j = 1; j < N - 1; j++) {
      for (int k = N >> 1; k > (i ^= k); k >>= 1)
        ;
      if (j < i)
        swap(a[i], a[j]);
    }
  }
  ll extgcd(ll a, ll b, ll &x, ll &y) {
    ll d = a;
    if (b != 0) {
      d = extgcd(b, a % b, y, x);
      y -= (a / b) * x;
    } else {
      x = 1;
      y = 0;
    }
    return d;
  }
  ll mod_inv(ll a) {
    ll x, y;
    extgcd(a, mod, x, y);
    return (mod + x % mod) % mod;
  }
  ll mod_pow(ll x, ll n) {
    if (n == 0)
      return 1;
    ll res = mod_pow(x * x % mod, n / 2);
    if (n & 1)
      res = res * x % mod;
    return res;
  }

  void _calc(vector<ll> &a, int sign) {
    int N = a.size();
    ll g = primitive_root;

    ll tmp = (mod - 1) * mod_inv(N) % mod;
    ll h = mod_pow(g, tmp);
    if (sign == -1)
      h = mod_inv(h);

    bit_reverse(a);

    for (int m = 1; m < N; m <<= 1) {
      ll _base = mod_pow(h, N / (2 * m));
      ll _w = 1;
      for (int x = 0; x < m; x++) {
        for (int s = x; s < N; s += (2 * m)) {
          ll u = a[s];
          ll d = (a[s + m] * _w) % mod;
          a[s] = (u + d) % mod;
          a[s + m] = (u - d + mod) % mod;
        }
        _w = (_w * _base) % mod;
      }
    }
  }

  void ntt(vector<ll> &in) { _calc(in, 1); }
  void intt(vector<ll> &in) {
    _calc(in, -1);
    ll n_inv = mod_inv(in.size());
    for (int i = 0; i < in.size(); i++) {
      in[i] = (in[i] * n_inv) % mod;
    }
  }

public:
  NTT() {}

  long long get_mod() const { return mod; }

  vector<ll> conv(const vector<ll> &a, const vector<ll> &b) {
    vector<ll> _a = a, _b = b;
    int m = a.size() + b.size() - 1;
    int n = 1;
    while (n < m)
      n <<= 1;
    _a.resize(n, 0);
    _b.resize(n, 0);
    ntt(_a);
    ntt(_b);
    for (int i = 0; i < n; i++)
      _a[i] = (_a[i] * _b[i]) % mod;
    intt(_a);
    _a.resize(m);
    return _a;
  }
};

NTT<998244353, 3> ntt;

vector<ll> dfs(const vector<int> &v, int S, int l, int r) {
  if (l + 1 == r) {
    vector<ll> x(min(v[l], S) + 1, 0);
    x[0] = 2;
    if (v[l] <= S)
      x[v[l]] = 1;
    return x;
  }

  int m = (l + r) / 2;
  vector<ll> lhs = dfs(v, S, l, m);
  vector<ll> rhs = dfs(v, S, m, r);

  vector<ll> res = ntt.conv(lhs, rhs);
  if (res.size() > S) {
    res.resize(S + 1);
  }
  return res;
}

int main() {
  int N, S;
  cin >> N >> S;

  vector<int> v;
  rep(i, N) {
    int a;
    cin >> a;
    v.push_back(a);
  }
  sort(ALLOF(v));

  vector<ll> ret = dfs(v, S, 0, N);

  cout << ret[S] << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, n) for (int i = (a); i < (int)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define FOR(it, c)                                                             \\
  for (__typeof((c).begin()) it = (c).begin(); it != (c).end(); ++it)
#define ALLOF(c) (c).begin(), (c).end()
typedef long long ll;
typedef unsigned long long ull;

template <ll mod, ll primitive_root> class NTT {
  void bit_reverse(vector<ll> &a) {
    int N = a.size();
    int i = 0;
    for (int j = 1; j < N - 1; j++) {
      for (int k = N >> 1; k > (i ^= k); k >>= 1)
        ;
      if (j < i)
        swap(a[i], a[j]);
    }
  }
  ll extgcd(ll a, ll b, ll &x, ll &y) {
    ll d = a;
    if (b != 0) {
      d = extgcd(b, a % b, y, x);
      y -= (a / b) * x;
    } else {
      x = 1;
      y = 0;
    }
    return d;
  }
  ll mod_inv(ll a) {
    ll x, y;
    extgcd(a, mod, x, y);
    return (mod + x % mod) % mod;
  }
  ll mod_pow(ll x, ll n) {
    if (n == 0)
      return 1;
    ll res = mod_pow(x * x % mod, n / 2);
    if (n & 1)
      res = res * x % mod;
    return res;
  }

  void _calc(vector<ll> &a, int sign) {
    int N = a.size();
    ll g = primitive_root;

    ll tmp = (mod - 1) * mod_inv(N) % mod;
    ll h = mod_pow(g, tmp);
    if (sign == -1)
      h = mod_inv(h);

    bit_reverse(a);

    for (int m = 1; m < N; m <<= 1) {
      ll _base = mod_pow(h, N / (2 * m));
      ll _w = 1;
      for (int x = 0; x < m; x++) {
        for (int s = x; s < N; s += (2 * m)) {
          ll u = a[s];
          ll d = (a[s + m] * _w) % mod;
          a[s] = (u + d) % mod;
          a[s + m] = (u - d + mod) % mod;
        }
        _w = (_w * _base) % mod;
      }
    }
  }

  void ntt(vector<ll> &in) { _calc(in, 1); }
  void intt(vector<ll> &in) {
    _calc(in, -1);
    ll n_inv = mod_inv(in.size());
    for (int i = 0; i < in.size(); i++) {
      in[i] = (in[i] * n_inv) % mod;
    }
  }

public:
  NTT() {}

  long long get_mod() const { return mod; }

  vector<ll> conv(const vector<ll> &a, const vector<ll> &b) {
    vector<ll> _a = a, _b = b;
    int m = a.size() + b.size() - 1;
    int n = 1;
    while (n < m)
      n <<= 1;
    _a.resize(n, 0);
    _b.resize(n, 0);
    ntt(_a);
    ntt(_b);
    for (int i = 0; i < n; i++)
      _a[i] = (_a[i] * _b[i]) % mod;
    intt(_a);
    _a.resize(m);
    return _a;
  }
};

NTT<998244353, 3> ntt;

vector<ll> dfs(const vector<int> &v, int S, int l, int r) {
  if (l + 1 == r) {
    vector<ll> x(S + 1, 0);
    x[0] = 2;
    if (v[l] <= S)
      x[v[l]] = 1;
    return x;
  }

  int m = (l + r) / 2;
  vector<ll> lhs = dfs(v, S, l, m);
  vector<ll> rhs = dfs(v, S, m, r);

  vector<ll> res = ntt.conv(lhs, rhs);
  if (res.size() > S) {
    res.resize(S + 1);
  }
  return res;
}

int main() {
  int N, S;
  cin >> N >> S;

  vector<int> v;
  rep(i, N) {
    int a;
    cin >> a;
    v.push_back(a);
  }

  vector<ll> ret = dfs(v, S, 0, N);

  cout << ret[S] << endl;

  return 0;
}