Problem p01749 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <unordered_map>
#include <unordered_set>
#include <valarray>
#include <vector>

using namespace std;
using ll = long long;
using ull = unsigned long long;
constexpr int TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n - 1); }

/// x^n
template <class T> T pow(T x, ll n) {
  T r = 1;
  while (n) {
    if (n & 1)
      r *= x;
    x *= x;
    n >>= 1;
  }
  return r;
}

template <uint MD> struct ModInt {
  uint v;
  ModInt() : v(0) {}
  ModInt(ll v) : v(normS(v % MD + MD)) {}
  uint value() const { return v; }
  static uint normS(const uint &x) { return (x < MD) ? x : x - MD; };
  static ModInt make(const uint &x) {
    ModInt m;
    m.v = x;
    return m;
  }
  const ModInt operator+(const ModInt &r) const { return make(normS(v + r.v)); }
  const ModInt operator-(const ModInt &r) const {
    return make(normS(v + normS(MD - r.v)));
  }
  const ModInt operator*(const ModInt &r) const {
    return make((ull)v * r.v % MD);
  }
  ModInt &operator+=(const ModInt &r) { return *this = *this + r; }
  ModInt &operator-=(const ModInt &r) { return *this = *this - r; }
  ModInt &operator*=(const ModInt &r) { return *this = *this * r; }
  static ModInt inv(const ModInt &x) { return pow(ModInt(x), MD - 2); }
};

struct Union {
  int B;
  vector<int> g;
  Union(int B) : B(B) {
    g.resize(B);
    iota(g.begin(), g.end(), 0);
  }
  void compress() {
    int co = 0;
    map<int, int> mp;
    for (int i = 0; i < B; i++) {
      if (!mp.count(g[i])) {
        mp[g[i]] = co++;
      }
      g[i] = mp[g[i]];
    }
  }
  void merge(int a, int b) {
    int c = g[b];
    for (int i = 0; i < B; i++) {
      if (g[i] == c)
        g[i] = g[a];
    }
  }
  int gc() {
    int ma = 0;
    for (int i = 0; i < B; i++) {
      ma = max(ma, g[i] + 1);
    }
    return ma;
  }
  const bool operator<(const Union &r) const {
    for (int i = 0; i < B; i++) {
      if (g[i] != r.g[i])
        return g[i] < r.g[i];
    }
    return false;
  }
  const bool operator==(const Union &r) const {
    for (int i = 0; i < B; i++) {
      if (g[i] != r.g[i])
        return false;
    }
    return true;
  }
  void print() {
    for (int i = 0; i < B; i++) {
      if (g[i] < 26) {
        printf("%c", g[i] + 'a');
      } else {
        printf("%c", g[i] + 'A' - 26);
      }
    }
    printf("\\n");
  }
};

const int MD = 1e9 + 7;
using Mint = ModInt<MD>;
const int MA = 210;
const int MB = 55;
const int MS = 2300;
int A, B, C;
vector<Mint> co;
vector<array<int, MB>> nex;

bool used[MA][MS][MB][2];
Mint dp[MA][MS][MB][2];
Mint solve(int P, int S, int K, int F) {
  K = min(K, B);
  if (P == 0) {
    assert(K == B);
    if (F)
      return 0;
    return co[S];
  }
  if (used[P][S][K][F])
    return dp[P][S][K][F];
  used[P][S][K][F] = true;
  Mint ans = solve(P - 1, S, K + 1, F);
  if (K == B)
    ans *= C;
  if (P >= B)
    ans += solve(P - 1, nex[S][K], 1, 1 - F);
  //    printf("solve %d %d %d %d %d\\n", P, S, K, F, ans.value());
  return dp[P][S][K][F] = ans;
}

int main() {
  cin >> A >> B >> C;

  Mint powC[MB];
  powC[0] = 1;
  for (int i = 1; i < MB; i++) {
    powC[i] = powC[i - 1] * C;
  }

  vector<Union> v;
  v.push_back(Union(B));

  for (int d = 1; d < B; d++) {
    vector<Union> v2(v.begin(), v.end());
    for (Union u : v) {
      for (int i = 0; i < B - d; i++) {
        u.merge(i, i + d);
      }
      u.compress();
      v2.push_back(u);
    }
    sort(v2.begin(), v2.end());
    v2.erase(unique(v2.begin(), v2.end()), v2.end());
    v = v2;
  }

  int m = (int)v.size();
  map<Union, int> u2id;
  for (int i = 0; i < m; i++) {
    u2id[v[i]] = i;
  }
  co.resize(m);
  nex.resize(m);
  for (int i = 0; i < m; i++) {
    co[i] = powC[v[i].gc()];
    for (int d = 1; d <= B; d++) {
      Union u = v[i];
      for (int j = 0; j < B - d; j++) {
        u.merge(j, j + d);
      }
      u.compress();
      nex[i][d] = u2id[u];
    }
  }

  /*    for (Union u: v) {
          u.print();
      }
      for (int i = 0; i < m; i++) {
          printf("%d co(%d) : ", i, co[i].value());
          for (int j = 0; j <= B; j++) {
              printf("%d ", nex[i][j]);
          }
          printf("\\n");
      }
      cout << "size:" << v.size() << endl;*/
  cout << (solve(A, m - 1, B, 1) - solve(A, m - 1, B, 0) +
           pow(Mint(C), B) * pow(Mint(C), A))
              .value()
       << endl;
  return 0;
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <unordered_map>
#include <unordered_set>
#include <valarray>
#include <vector>

using namespace std;
using ll = long long;
using ull = unsigned long long;
constexpr int TEN(int n) { return (n == 0) ? 1 : 10 * TEN(n - 1); }

/// x^n
template <class T> T pow(T x, ll n) {
  T r = 1;
  while (n) {
    if (n & 1)
      r *= x;
    x *= x;
    n >>= 1;
  }
  return r;
}

template <uint MD> struct ModInt {
  uint v;
  ModInt() : v(0) {}
  ModInt(ll v) : v(normS(v % MD + MD)) {}
  uint value() const { return v; }
  static uint normS(const uint &x) { return (x < MD) ? x : x - MD; };
  static ModInt make(const uint &x) {
    ModInt m;
    m.v = x;
    return m;
  }
  const ModInt operator+(const ModInt &r) const { return make(normS(v + r.v)); }
  const ModInt operator-(const ModInt &r) const {
    return make(normS(v + normS(MD - r.v)));
  }
  const ModInt operator*(const ModInt &r) const {
    return make((ull)v * r.v % MD);
  }
  ModInt &operator+=(const ModInt &r) { return *this = *this + r; }
  ModInt &operator-=(const ModInt &r) { return *this = *this - r; }
  ModInt &operator*=(const ModInt &r) { return *this = *this * r; }
  static ModInt inv(const ModInt &x) { return pow(ModInt(x), MD - 2); }
};

struct Union {
  int B;
  vector<int> g;
  Union(int B) : B(B) {
    g.resize(B);
    iota(g.begin(), g.end(), 0);
  }
  void compress() {
    int co = 0;
    map<int, int> mp;
    for (int i = 0; i < B; i++) {
      if (!mp.count(g[i])) {
        mp[g[i]] = co++;
      }
      g[i] = mp[g[i]];
    }
  }
  void merge(int a, int b) {
    int c = g[b];
    for (int i = 0; i < B; i++) {
      if (g[i] == c)
        g[i] = g[a];
    }
  }
  int gc() {
    int ma = 0;
    for (int i = 0; i < B; i++) {
      ma = max(ma, g[i] + 1);
    }
    return ma;
  }
  const bool operator<(const Union &r) const {
    for (int i = 0; i < B; i++) {
      if (g[i] != r.g[i])
        return g[i] < r.g[i];
    }
    return false;
  }
  const bool operator==(const Union &r) const {
    for (int i = 0; i < B; i++) {
      if (g[i] != r.g[i])
        return false;
    }
    return true;
  }
  void print() {
    for (int i = 0; i < B; i++) {
      if (g[i] < 26) {
        printf("%c", g[i] + 'a');
      } else {
        printf("%c", g[i] + 'A' - 26);
      }
    }
    printf("\\n");
  }
};

const int MD = 1e9 + 7;
using Mint = ModInt<MD>;
const int MA = 1100;
const int MB = 55;
const int MS = 2300;
int A, B, C;
vector<Mint> co;
vector<array<int, MB>> nex;

bool used[MA][MS][MB][2];
Mint dp[MA][MS][MB][2];
Mint solve(int P, int S, int K, int F) {
  K = min(K, B);
  if (P == 0) {
    assert(K == B);
    if (F)
      return 0;
    return co[S];
  }
  if (used[P][S][K][F])
    return dp[P][S][K][F];
  used[P][S][K][F] = true;
  Mint ans = solve(P - 1, S, K + 1, F);
  if (K == B)
    ans *= C;
  if (P >= B)
    ans += solve(P - 1, nex[S][K], 1, 1 - F);
  //    printf("solve %d %d %d %d %d\\n", P, S, K, F, ans.value());
  return dp[P][S][K][F] = ans;
}

int main() {
  cin >> A >> B >> C;

  Mint powC[MB];
  powC[0] = 1;
  for (int i = 1; i < MB; i++) {
    powC[i] = powC[i - 1] * C;
  }

  vector<Union> v;
  v.push_back(Union(B));

  for (int d = 1; d < B; d++) {
    vector<Union> v2(v.begin(), v.end());
    for (Union u : v) {
      for (int i = 0; i < B - d; i++) {
        u.merge(i, i + d);
      }
      u.compress();
      v2.push_back(u);
    }
    sort(v2.begin(), v2.end());
    v2.erase(unique(v2.begin(), v2.end()), v2.end());
    v = v2;
  }

  int m = (int)v.size();
  map<Union, int> u2id;
  for (int i = 0; i < m; i++) {
    u2id[v[i]] = i;
  }
  co.resize(m);
  nex.resize(m);
  for (int i = 0; i < m; i++) {
    co[i] = powC[v[i].gc()];
    for (int d = 1; d <= B; d++) {
      Union u = v[i];
      for (int j = 0; j < B - d; j++) {
        u.merge(j, j + d);
      }
      u.compress();
      nex[i][d] = u2id[u];
    }
  }

  /*    for (Union u: v) {
          u.print();
      }
      for (int i = 0; i < m; i++) {
          printf("%d co(%d) : ", i, co[i].value());
          for (int j = 0; j <= B; j++) {
              printf("%d ", nex[i][j]);
          }
          printf("\\n");
      }
      cout << "size:" << v.size() << endl;*/
  cout << (solve(A, m - 1, B, 1) - solve(A, m - 1, B, 0) +
           pow(Mint(C), B) * pow(Mint(C), A))
              .value()
       << endl;
  return 0;
}