Problem p03322 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <bitset>
#include <cassert>
#include <ciso646>
#include <cmath>
#include <complex>
#include <cstdio>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std;

// #define int long long
typedef long long ll;

typedef unsigned long long ul;
typedef unsigned int ui;
// const ll mod = 1000000087;
// const ll INF = mod * mod;
typedef pair<int, int> P;
typedef pair<int, bool> sP;
#define stop                                                                   \
  char nyaa;                                                                   \
  cin >> nyaa;
#define rep(i, n) for (int i = 0; i < n; i++)
#define per(i, n) for (int i = n - 1; i >= 0; i--)
#define Rep(i, sta, n) for (int i = sta; i < n; i++)
#define rep1(i, n) for (int i = 1; i <= n; i++)
#define per1(i, n) for (int i = n; i >= 1; i--)
#define Rep1(i, sta, n) for (int i = sta; i <= n; i++)
typedef pair<ll, ll> LP;
typedef vector<ll> vec;
typedef long double ld;
typedef pair<ld, ld> LDP;
const ld eps = 1e-8;
const ld pi = acos(-1.0);
typedef vector<vector<ll>> mat;
typedef vector<ll> vec;

vector<ll> mods;
vector<ll> roots;

bool isp(ll x) {
  Rep1(i, 2, 100000) {
    if (x % i == 0)
      return false;
  }
  return true;
}
void genmods() {
  ll sta = 1000010000;
  while (mods.size() < 3) {
    if (isp(sta))
      mods.push_back(sta);
    sta++;
  }
}
void genroots() {
  ll sta = 2000010000;
  while (roots.size() < 2) {
    if (isp(sta))
      roots.push_back(sta);
    sta++;
  }
}
ll mod_pow(ll x, ll n, ll mod) {
  ll res = 1;
  while (n) {
    if (n % 2)
      res = res * x % mod;
    x = x * x % mod;
    n >>= 1;
  }
  return res;
}
vector<ll> invroots;

vector<ll> tr[300000];
vector<ll> invtr[300000];
void init() {
  genmods();
  genroots();
  rep(i, 300000) {
    tr[i].resize(mods.size() * roots.size());
    invtr[i].resize(mods.size() * roots.size());
  }
  rep(j, tr[0].size()) { tr[0][j] = invtr[0][j] = 1; }
  rep(j, mods.size()) {
    rep(k, roots.size()) {
      invroots.push_back(mod_pow(roots[k], mods[j] - 2, mods[j]));
    }
  }
  rep(i, 299999) {
    rep(j, mods.size()) {
      rep(k, roots.size()) {
        int id = j * roots.size() + k;
        tr[i + 1][id] = tr[i][id] * roots[k] % mods[j];
        invtr[i + 1][id] = invtr[i][id] * invroots[id] % mods[j];
      }
    }
  }
}

void solve() {
  int n;
  cin >> n;
  string s;
  cin >> s;
  vector<ll> sum(mods.size() * roots.size(), 0);
  int loc = 0;
  vector<vector<ll>> memo(n);
  rep(i, n) memo[i].resize(mods.size() * roots.size());
  vector<int> locs(n);
  int len = roots.size() * mods.size();
  rep(i, n) {
    if (s[i] == '>')
      loc++;
    else if (s[i] == '<')
      loc--;
    else if (s[i] == '+') {
      if (loc >= 0) {
        rep(j, len) { sum[j] += tr[loc][j]; }
      } else {
        rep(j, len) { sum[j] += invtr[-loc][j]; }
      }
    } else {
      if (loc >= 0) {
        rep(j, len) { sum[j] -= tr[loc][j]; }
      } else {
        rep(j, len) { sum[j] -= invtr[-loc][j]; }
      }
    }
    rep(j, len) {
      ll mod = mods[j / roots.size()];
      if (sum[j] < 0)
        sum[j] += mod;
      if (sum[j] >= mod)
        sum[j] -= mod;
      memo[i][j] = sum[j];
    }

    locs[i] = loc;
  }
  vector<ll> alls = sum;
  ll ans = 0;
  vector<map<ll, int>> mp(len);
  rep(i, len) { mp[i][alls[i]]++; }
  rep(i, n) {
    // cout << mp[memo[i]] << endl;
    int mi = 2 * n;
    rep(j, len) { mi = min(mi, mp[j][memo[i][j]]); }
    ans += mi;
    bool valid = true;

    rep(j, len) {
      ll nex = alls[j];
      ll mod = mods[j / roots.size()];
      if (locs[i] >= 0) {
        nex = nex * tr[locs[i]][j] % mod;
      } else {
        nex = nex * invtr[-locs[i]][j] % mod;
      }
      nex += memo[i][j];
      if (nex >= mod)
        nex -= mod;
      mp[j][nex]++;
    }
  }
  cout << ans << endl;
}

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  // cout << fixed << setprecision(7);
  init();
  solve();
  // stop
  return 0;
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cassert>
#include <ciso646>
#include <cmath>
#include <complex>
#include <cstdio>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace std;

// #define int long long
typedef long long ll;

typedef unsigned long long ul;
typedef unsigned int ui;
// const ll mod = 1000000087;
// const ll INF = mod * mod;
typedef pair<int, int> P;
typedef pair<int, bool> sP;
#define stop                                                                   \
  char nyaa;                                                                   \
  cin >> nyaa;
#define rep(i, n) for (int i = 0; i < n; i++)
#define per(i, n) for (int i = n - 1; i >= 0; i--)
#define Rep(i, sta, n) for (int i = sta; i < n; i++)
#define rep1(i, n) for (int i = 1; i <= n; i++)
#define per1(i, n) for (int i = n; i >= 1; i--)
#define Rep1(i, sta, n) for (int i = sta; i <= n; i++)
typedef pair<ll, ll> LP;
typedef vector<ll> vec;
typedef long double ld;
typedef pair<ld, ld> LDP;
const ld eps = 1e-8;
const ld pi = acos(-1.0);
typedef vector<vector<ll>> mat;
typedef vector<ll> vec;

vector<ll> mods;
vector<ll> roots;

bool isp(ll x) {
  Rep1(i, 2, 100000) {
    if (x % i == 0)
      return false;
  }
  return true;
}
void genmods() {
  ll sta = 1000010000;
  while (mods.size() < 3) {
    if (isp(sta))
      mods.push_back(sta);
    sta++;
  }
}
void genroots() {
  ll sta = 2000010000;
  while (roots.size() < 2) {
    if (isp(sta))
      roots.push_back(sta);
    sta++;
  }
}
ll mod_pow(ll x, ll n, ll mod) {
  ll res = 1;
  while (n) {
    if (n % 2)
      res = res * x % mod;
    x = x * x % mod;
    n >>= 1;
  }
  return res;
}
vector<ll> invroots;

vector<ll> tr[300000];
vector<ll> invtr[300000];
void init() {
  genmods();
  genroots();
  rep(i, 300000) {
    tr[i].resize(mods.size() * roots.size());
    invtr[i].resize(mods.size() * roots.size());
  }
  rep(j, tr[0].size()) { tr[0][j] = invtr[0][j] = 1; }
  rep(j, mods.size()) {
    rep(k, roots.size()) {
      invroots.push_back(mod_pow(roots[k], mods[j] - 2, mods[j]));
    }
  }
  rep(i, 299999) {
    rep(j, mods.size()) {
      rep(k, roots.size()) {
        int id = j * roots.size() + k;
        tr[i + 1][id] = tr[i][id] * roots[k] % mods[j];
        invtr[i + 1][id] = invtr[i][id] * invroots[id] % mods[j];
      }
    }
  }
}

void solve() {
  int n;
  cin >> n;
  string s;
  cin >> s;
  vector<ll> sum(mods.size() * roots.size(), 0);
  int loc = 0;
  vector<vector<ll>> memo(n);
  rep(i, n) memo[i].resize(mods.size() * roots.size());
  vector<int> locs(n);
  int len = roots.size() * mods.size();
  rep(i, n) {
    if (s[i] == '>')
      loc++;
    else if (s[i] == '<')
      loc--;
    else if (s[i] == '+') {
      if (loc >= 0) {
        rep(j, len) { sum[j] += tr[loc][j]; }
      } else {
        rep(j, len) { sum[j] += invtr[-loc][j]; }
      }
    } else {
      if (loc >= 0) {
        rep(j, len) { sum[j] -= tr[loc][j]; }
      } else {
        rep(j, len) { sum[j] -= invtr[-loc][j]; }
      }
    }
    rep(j, len) {
      ll mod = mods[j / roots.size()];
      if (sum[j] < 0)
        sum[j] += mod;
      if (sum[j] >= mod)
        sum[j] -= mod;
      memo[i][j] = sum[j];
    }

    locs[i] = loc;
  }
  vector<ll> alls = sum;
  ll ans = 0;
  vector<unordered_map<ll, int>> mp(len);
  rep(i, len) { mp[i][alls[i]]++; }
  rep(i, n) {
    // cout << mp[memo[i]] << endl;
    int mi = 2 * n;
    rep(j, len) { mi = min(mi, mp[j][memo[i][j]]); }
    ans += mi;
    bool valid = true;

    rep(j, len) {
      ll nex = alls[j];
      ll mod = mods[j / roots.size()];
      if (locs[i] >= 0) {
        nex = nex * tr[locs[i]][j] % mod;
      } else {
        nex = nex * invtr[-locs[i]][j] % mod;
      }
      nex += memo[i][j];
      if (nex >= mod)
        nex -= mod;
      mp[j][nex]++;
    }
  }
  cout << ans << endl;
}

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  // cout << fixed << setprecision(7);
  init();
  solve();
  // stop
  return 0;
}