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

#include <algorithm>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstring>
#include <iostream>
#include <map>
#include <random>
#include <unordered_map>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;

ll gcd(ll a, ll b) {
  if (!a)
    return b;
  return gcd(b % a, a);
}

const int MN = 1010;
const ll MD = 1e9 + 7;

ll m;

ll inv[MN];
ll powt[MN][MN];
ll fact[MN];
ll comb[MN][MN];
ll gcdm[MN];
ll get2[MN];
void init() {
  inv[1] = 1;
  for (int i = 2; i < MN; i++) {
    inv[i] = inv[MD % i] * (MD - MD / i) % MD;
  }
  for (int i = 0; i < MN; i++) {
    powt[i][0] = 1;
    for (int j = 1; j < MN; j++) {
      powt[i][j] = powt[i][j - 1] * i % MD;
    }
  }
  fact[0] = 1;
  for (int i = 1; i < MN; i++) {
    fact[i] = fact[i - 1] * i % MD;
  }
  comb[0][0] = 1;
  for (int i = 1; i < MN; i++) {
    comb[i][0] = 1;
    for (int j = 1; j < MN; j++) {
      comb[i][j] = comb[i - 1][j] + comb[i - 1][j - 1];
      comb[i][j] %= MD;
    }
  }
  for (int i = 0; i < MN; i++) {
    gcdm[i] = gcd(m, i);
  }
  get2[0] = 1;
  get2[1] = 1;
  for (int i = 2; i < MN; i++) {
    get2[i] = get2[i - 1] + (i - 1) * get2[i - 2];
    get2[i] %= MD;
  }
}

ll dp[3][MN][MN];
ll calc(int f, int n, int k) {
  if (k < 0)
    return 0;
  if (!k)
    return 1;
  if (dp[f][n][k] != -1)
    return dp[f][n][k];
  if (!f) {
    ll res = 0;
    for (int i = 1; i <= k; i++) {
      if (n * i / gcdm[n * i] != n)
        continue;
      ll r = calc(f, n, k - i);
      r *= comb[n * k - 1][n * i - 1];
      r %= MD;
      r *= fact[n * i - 1];
      r %= MD;
      res += r;
      res %= MD;
    }
    return (dp[f][n][k] = res);
  } else {
    ll res = (n % 2) * calc(f, n, k - 1);
    res %= MD;
    ll r = (k - 1) * n;
    r %= MD;
    r *= calc(f, n, k - 2);
    r %= MD;
    res += r;
    res %= MD;
    return (dp[f][n][k] = res);
  }
}
ll solve(int n, int m) {
  assert(0 <= n);
  if (!n)
    return 1;
  if (n < m)
    return 0;
  if (dp[2][n][m] != -1)
    return dp[2][n][m];
  ll res = solve(n, m + 1);
  for (int i = 1; i <= n / m; i++) {
    ll r = comb[n][i * m];
    r *= calc(0, m, i);
    r %= MD;
    r *= calc(1, m, i);
    r %= MD;
    r *= solve(n - i * m, m + 1);
    r %= MD;
    res += r;
    res %= MD;
  }
  return (dp[2][n][m] = res);
}

int main() {
  memset(dp, -1, sizeof(dp));
  ll n, x, y, z;
  cin >> n >> x >> y >> z;
  x += z;
  m = abs(x - y);
  init();
  if (!x && !y) {
    cout << powt[n][n] * get2[n] % MD << endl;
    return 0;
  }
  if (x == y) {
    cout << fact[n] * get2[n] % MD << endl;
    return 0;
  }
  cout << solve(n, 1) << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cmath>
#include <complex>
#include <iostream>
#include <map>
#include <random>
#include <unordered_map>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;

const int MN = 1010;
const ll MD = 1e9 + 7;

struct Perm {
  int n;
  vector<int> d;
  Perm(int n) : n(n) {
    d.resize(n);
    for (int i = 0; i < n; i++) {
      d[i] = i;
    }
  }

  int &operator[](int p) { return d[p]; }

  const int &operator[](int p) const { return d[p]; }

  Perm operator*(const Perm &r) {
    assert(n == r.n);
    Perm res(n);
    for (int i = 0; i < n; i++) {
      res[i] = r[d[i]];
    }
    return res;
  }

  Perm pow(ll p) {
    Perm res(n), buf = *this;
    while (p != 0) {
      if (p % 2) {
        res = res * buf;
      }
      p /= 2;
      buf = buf * buf;
    }
    return res;
  }
};

int main() {
  int n, x, y, z;
  cin >> n >> x >> y >> z;
  int res = 0;
  vector<int> b1(n), b2(n);
  for (int i = 0; i < n; i++) {
    b1[i] = i;
  }
  do {
    int r = 0;
    for (int i = 0; i < n; i++) {
      b2[i] = i;
    }
    do {
      Perm p1(n), p2(n);
      p1.d = b1;
      p2.d = b2;
      Perm p = p1.pow(x) * p2 * p1.pow(y) * p2 * p1.pow(z);
      bool f = true;
      for (int i = 0; i < n; i++) {
        if (p[i] != i) {
          f = false;
          break;
        }
      }
      if (f) {
        vector<int> u = {0, 2, 3, 1};
        if (b1 == u) {
          for (int i = 0; i < n; i++) {
            cout << b2[i] << " ";
          }
          cout << endl;
        }
        r++;
      }
    } while (next_permutation(b2.begin(), b2.end()));
    for (int i = 0; i < n; i++) {
      cout << b1[i] << " ";
    }
    cout << "res: " << r << endl;
    res += r;
  } while (next_permutation(b1.begin(), b1.end()));
  cout << res << endl;
  return 0;
}