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

#include <bits/stdc++.h>
using namespace std;
using Int = signed;
using ll = long long;

Int extgcd(Int a, Int b, Int &x, Int &y) {
  Int d = a;
  if (b != 0) {
    d = extgcd(b, a % b, y, x);
    y -= (a / b) * x;
  } else {
    x = 1;
    y = 0;
  }
  return d;
}

Int mod_inverse(Int a, Int mod) {
  Int x, y;
  extgcd(a, mod, x, y);
  return (mod + x % mod) % mod;
}

using V = tuple<Int, Int, Int>;
using T = long long;
const Int MOD1 = 1e9 + 7;
const Int MOD2 = 1e9 + 9;
const Int MOD3 = 1e8 + 7;

inline T add1(T a, T b) {
  a += b;
  if (a >= MOD1)
    a -= MOD1;
  return a;
}

inline T add2(T a, T b) {
  a += b;
  if (a >= MOD2)
    a -= MOD2;
  return a;
}

inline T add3(T a, T b) {
  a += b;
  if (a >= MOD3)
    a -= MOD3;
  return a;
}

inline T sub1(T a, T b) { return add1(a, MOD1 - b); }

inline T sub2(T a, T b) { return add2(a, MOD2 - b); }

inline T sub3(T a, T b) { return add3(a, MOD3 - b); }

inline T mul1(T a, T b) { return (1LL * a * b) % MOD1; }

inline T mul2(T a, T b) { return (1LL * a * b) % MOD2; }

inline T mul3(T a, T b) { return (1LL * a * b) % MOD3; }

T inv1(T a) { return mod_inverse(a, MOD1); }
T inv2(T a) { return mod_inverse(a, MOD2); }
T inv3(T a) { return mod_inverse(a, MOD3); }

V inv(const V &a) {
  T a1, a2, a3;
  tie(a1, a2, a3) = a;
  T r1 = inv1(a1);
  T r2 = inv2(a2);
  T r3 = inv3(a3);
  return V(r1, r2, r3);
}

V add(const V &a, const V &b) {
  T a1, a2, a3, b1, b2, b3;
  tie(a1, a2, a3) = a;
  tie(b1, b2, b3) = b;
  T r1 = add1(a1, b1);
  T r2 = add2(a2, b2);
  T r3 = add3(a3, b3);
  return V(r1, r2, r3);
}

V sub(const V &a, const V &b) {
  T a1, a2, a3, b1, b2, b3;
  tie(a1, a2, a3) = a;
  tie(b1, b2, b3) = b;
  T r1 = sub1(a1, b1);
  T r2 = sub2(a2, b2);
  T r3 = sub3(a3, b3);
  return V(r1, r2, r3);
}

V mul(const V &a, const V &b) {
  T a1, a2, a3, b1, b2, b3;
  tie(a1, a2, a3) = a;
  tie(b1, b2, b3) = b;
  T r1 = mul1(a1, b1);
  T r2 = mul2(a2, b2);
  T r3 = mul3(a3, b3);
  return V(r1, r2, r3);
}

vector<vector<V>> vs;
vector<vector<vector<V>>> vm;

V expr(string &s, Int &p, Int d);
V term(string &s, Int &p, Int d);
V factor(string &s, Int &p, Int d);
V number(string &s, Int &p);

V expr(string &s, Int &p, Int d) {
  V res = term(s, p, d);
  vs[d].emplace_back(res);
  while (p < (Int)s.size()) {
    if (s[p] == '+') {
      p++;
      V b = term(s, p, d);
      vs[d].emplace_back(b);
      res = add(res, b);
      continue;
    }
    break;
  }
  vs[d].emplace_back(V(0, 0, 0));
  vm[d].emplace_back();
  return res;
}
V term(string &s, Int &p, Int d) {
  V res = factor(s, p, d);
  vm[d].emplace_back();
  vm[d].back().emplace_back(res);
  while (p < (Int)s.size()) {
    if (s[p] == '*') {
      p++;
      V b = factor(s, p, d);
      vm[d].back().emplace_back(b);
      res = mul(res, b);
      continue;
    }
    break;
  }
  return res;
}
V factor(string &s, Int &p, Int d) {
  if (s[p] == '(') {
    p++;
    V res = expr(s, p, d + 1);
    assert(s[p] == ')');
    p++;
    return res;
  }
  return number(s, p);
}

V number(string &s, Int &p) {
  Int k = s[p++] - '0';
  return V(k, k, k);
}

signed main() {
  Int n;
  while (cin >> n, n) {
    string s;
    cin >> s;
    vs.assign(1010, vector<V>());
    vm.assign(1010, vector<vector<V>>());

    V nv(n % MOD1, n % MOD2, n % MOD3);
    Int p = 0;
    expr(s, p, 0);
    ll ans = 0;

    if (0) {
      for (Int i = 0; i < 3; i++) {
        for (Int j = 0; j < (Int)vs[i].size(); j++)
          cout << get<0>(vs[i][j]) << " ";
        cout << endl;
        for (Int j = 0; j < (Int)vm[i].size(); j++) {
          cout << i << " " << j << ":" << endl;
          for (Int k = 0; k < (Int)vm[i][j].size(); k++)
            cout << get<0>(vm[i][j][k]) << " ";
          cout << endl;
        }
        cout << endl;
      }
    }

    for (Int i = 0; i < 1010; i++) {
      auto &v = vs[i];
      if (v.empty())
        continue;
      map<V, ll> dp;
      V uku(0, 0, 0);
      dp[uku]++;
      for (Int j = 0; j < (Int)v.size(); j++) {
        // cout<<i<<" "<<j<<":"<<get<0>(v[j])<<endl;
        if (v[j] == V(0, 0, 0)) {
          dp.clear();
          uku = V(0, 0, 0);
          dp[uku]++;
          continue;
        }
        V malta(1, 1, 1);
        for (Int k = 0; k < (Int)vm[i][j].size(); k++) {
          malta = mul(malta, vm[i][j][k]);
          ans += dp[sub(add(uku, malta), nv)];
        }

        map<V, ll> dp2;
        malta = V(1, 1, 1);
        dp2[mul(malta, nv)]++;
        for (Int k = 1; k < (Int)vm[i][j].size(); k++) {
          malta = mul(malta, vm[i][j][k]);
          ans += dp2[malta];
          dp2[mul(malta, nv)]++;
        }

        uku = add(uku, v[j]);
        malta = V(1, 1, 1);
        for (Int k = 0; k + 1 < (Int)vm[i][j].size(); k++) {
          malta = mul(malta, vm[i][j][k]);
          dp[sub(uku, mul(v[j], inv(malta)))]++;
        }

        dp[uku]++;
      }
    }

    cout << ans << endl;
  }
  return 0;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using Int = long long;

Int extgcd(Int a, Int b, Int &x, Int &y) {
  Int d = a;
  if (b != 0) {
    d = extgcd(b, a % b, y, x);
    y -= (a / b) * x;
  } else {
    x = 1;
    y = 0;
  }
  return d;
}

Int mod_inverse(Int a, Int mod) {
  Int x, y;
  extgcd(a, mod, x, y);
  return (mod + x % mod) % mod;
}

using V = tuple<Int, Int, Int>;
using T = long long;
const Int MOD1 = 1e9 + 7;
const Int MOD2 = 1e9 + 9;
const Int MOD3 = 1e8 + 7;

inline T add1(T a, T b) {
  a += b;
  if (a >= MOD1)
    a -= MOD1;
  return a;
}

inline T add2(T a, T b) {
  a += b;
  if (a >= MOD2)
    a -= MOD2;
  return a;
}

inline T add3(T a, T b) {
  a += b;
  if (a >= MOD3)
    a -= MOD3;
  return a;
}

inline T sub1(T a, T b) { return add1(a, MOD1 - b); }

inline T sub2(T a, T b) { return add2(a, MOD2 - b); }

inline T sub3(T a, T b) { return add3(a, MOD3 - b); }

inline T mul1(T a, T b) { return (1LL * a * b) % MOD1; }

inline T mul2(T a, T b) { return (1LL * a * b) % MOD2; }

inline T mul3(T a, T b) { return (1LL * a * b) % MOD3; }

T inv1(T a) { return mod_inverse(a, MOD1); }
T inv2(T a) { return mod_inverse(a, MOD2); }
T inv3(T a) { return mod_inverse(a, MOD3); }

V inv(const V &a) {
  T a1, a2, a3;
  tie(a1, a2, a3) = a;
  T r1 = inv1(a1);
  T r2 = inv2(a2);
  T r3 = inv3(a3);
  return V(r1, r2, r3);
}

V add(const V &a, const V &b) {
  T a1, a2, a3, b1, b2, b3;
  tie(a1, a2, a3) = a;
  tie(b1, b2, b3) = b;
  T r1 = add1(a1, b1);
  T r2 = add2(a2, b2);
  T r3 = add3(a3, b3);
  return V(r1, r2, r3);
}

V sub(const V &a, const V &b) {
  T a1, a2, a3, b1, b2, b3;
  tie(a1, a2, a3) = a;
  tie(b1, b2, b3) = b;
  T r1 = sub1(a1, b1);
  T r2 = sub2(a2, b2);
  T r3 = sub3(a3, b3);
  return V(r1, r2, r3);
}

V mul(const V &a, const V &b) {
  T a1, a2, a3, b1, b2, b3;
  tie(a1, a2, a3) = a;
  tie(b1, b2, b3) = b;
  T r1 = mul1(a1, b1);
  T r2 = mul2(a2, b2);
  T r3 = mul3(a3, b3);
  return V(r1, r2, r3);
}

vector<vector<V>> vs;
vector<vector<vector<V>>> vm;

V expr(string &s, Int &p, Int d);
V term(string &s, Int &p, Int d);
V factor(string &s, Int &p, Int d);
V number(string &s, Int &p);

V expr(string &s, Int &p, Int d) {
  V res = term(s, p, d);
  vs[d].emplace_back(res);
  while (p < (Int)s.size()) {
    if (s[p] == '+') {
      p++;
      V b = term(s, p, d);
      vs[d].emplace_back(b);
      res = add(res, b);
      continue;
    }
    break;
  }
  return res;
}
V term(string &s, Int &p, Int d) {
  V res = factor(s, p, d);
  vm[d].emplace_back();
  vm[d].back().emplace_back(res);
  while (p < (Int)s.size()) {
    if (s[p] == '*') {
      p++;
      V b = factor(s, p, d);
      vm[d].back().emplace_back(b);
      res = mul(res, b);
      continue;
    }
    break;
  }
  return res;
}
V factor(string &s, Int &p, Int d) {
  if (s[p] == '(') {
    p++;
    V res = expr(s, p, d + 1);
    assert(s[p] == ')');
    p++;
    return res;
  }
  return number(s, p);
}

V number(string &s, Int &p) {
  Int k = s[p++] - '0';
  return V(k, k, k);
}

signed main() {
  Int n;
  while (cin >> n, n) {
    string s;
    cin >> s;
    vs.assign(1010, vector<V>());
    vm.assign(1010, vector<vector<V>>());

    V nv(n % MOD1, n % MOD2, n % MOD3);
    Int p = 0;
    expr(s, p, 0);
    Int ans = 0;

    if (0) {
      for (Int i = 0; i < 3; i++) {
        for (Int j = 0; j < (Int)vs[i].size(); j++)
          cout << get<0>(vs[i][j]) << " ";
        cout << endl;
        for (Int j = 0; j < (Int)vm[i].size(); j++) {
          cout << i << " " << j << ":" << endl;
          for (Int k = 0; k < (Int)vm[i][j].size(); k++)
            cout << get<0>(vm[i][j][k]) << " ";
          cout << endl;
        }
        cout << endl;
      }
    }

    for (Int i = 0; i < 1010; i++) {
      auto &v = vs[i];
      if (v.empty())
        continue;
      map<V, Int> dp;
      V uku(0, 0, 0);
      dp[uku]++;
      for (Int j = 0; j < (Int)v.size(); j++) {
        V malta(1, 1, 1);
        for (Int k = 0; k < (Int)vm[i][j].size(); k++) {
          malta = mul(malta, vm[i][j][k]);
          ans += dp[sub(add(uku, malta), nv)];
        }
        malta = V(1, 1, 1);
        for (Int k = 0; k + 1 < (Int)vm[i][j].size(); k++) {
          malta = mul(malta, vm[i][j][k]);
          dp[sub(add(uku, v[j]), mul(v[j], inv(malta)))]++;
        }
        malta = V(1, 1, 1);
        map<V, Int> dp2;
        dp2[mul(malta, nv)]++;
        for (Int k = 1; k < (Int)vm[i][j].size(); k++) {
          malta = mul(malta, vm[i][j][k]);
          ans += dp2[malta];
          dp2[mul(malta, nv)]++;
        }

        uku = add(uku, v[j]);
        dp[uku]++;
        // if(v[j]==nv) ans++;
      }
    }

    cout << ans << endl;
  }
  return 0;
}