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

#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long ull;
struct rollinghash {
  string S;
  ull B;
  vector<ull> hash, p;
  int len;
  rollinghash() {}
  rollinghash(string S, ull B) : S(S), B(B) { init(); };
  void init() {
    len = S.length();
    hash.resize(len + 1);
    p.resize(len + 1);
    hash[0] = 0;
    p[0] = 1;
    for (int i = 0; i < len; i++) {
      hash[i + 1] = hash[i] * B + S[i];
      p[i + 1] = p[i] * B;
    }
  }
  ull find(int l, int r) { return hash[r] - hash[l] * p[r - l]; }
};
typedef pair<ull, ull> P;
int main() {
  string s, t;
  cin >> s >> t;
  ull B = 1000000007LL, C = 1777771LL;
  int sl = s.length(), tl = t.length();
  rollinghash r1(s, B), r2(s, C);
  vector<P> v(sl - tl + 1);
  for (int i = 0; i <= sl - tl; i++)
    v[i] = P(r1.find(i, i + tl), r2.find(i, i + tl));
  sort(v.begin(), v.end());
  // for(int i=0;i<=sl-tl;i++) cout<<v[i].first<<" "<<v[i].second<<endl;
  // cout<<endl;
  ull h1 = 0, p1 = 1, h2 = 0, p2 = 1, ans = 0;
  for (int i = 0; i < tl; i++)
    h1 = h1 * B + t[i], h2 = h2 * C + t[i];
  for (int i = tl - 1; i >= 0; i--) {
    h1 -= p1 * t[i];
    h2 -= p2 * t[i];
    for (int j = 'a'; j <= 'z'; j++) {
      if (j == t[i])
        continue;
      h1 += p1 * j;
      h2 += p2 * j;
      // cout<<h1<<endl;
      ans += upper_bound(v.begin(), v.end(), P(h1, h2)) -
             lower_bound(v.begin(), v.end(), P(h1, h2));
      h1 -= p1 * j;
      h2 -= p2 * j;
    }
    for (int j = 'A'; j <= 'Z'; j++) {
      if (j == t[i])
        continue;
      h1 += p1 * j;
      h2 += p2 * j;
      // cout<<h1<<endl;
      ans += upper_bound(v.begin(), v.end(), P(h1, h2)) -
             lower_bound(v.begin(), v.end(), P(h1, h2));
      h1 -= p1 * j;
      h2 -= p2 * j;
    }
    h1 += p1 * t[i];
    h2 += p2 * t[i];
    p1 *= B;
    p2 *= C;
  }
  cout << ans << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define MAX_N 400050
int n, k;
int r[MAX_N + 1], r2[MAX_N + 1];
int t[MAX_N + 1];
bool compare_sa(int i, int j) {
  if (r[i] != r[j])
    return r[i] < r[j];
  else {
    int ri = i + k <= n ? r[i + k] : -1;
    int rj = j + k <= n ? r[j + k] : -1;
    return ri < rj;
  }
}
void constract_sa(const string &S, int *sa) {
  n = S.length();
  for (int i = 0; i <= n; i++) {
    sa[i] = i;
    r[i] = i < n ? S[i] : -1;
  }
  for (k = 1; k <= n; k *= 2) {
    sort(sa, sa + n + 1, compare_sa);
    t[sa[0]] = 0;
    for (int i = 0; i <= n; i++) {
      t[sa[i]] = t[sa[i - 1]] + (compare_sa(sa[i - 1], sa[i]) ? 1 : 0);
    }
    for (int i = 0; i <= n; i++) {
      r[i] = t[i];
    }
  }
}

int contains(const string &S, int *sa, const string &T) {
  int a = 0, b = S.length();
  while (b - a > 1) {
    int c = (a + b) / 2;
    if (S.compare(sa[c], T.length(), T) < 0)
      a = c;
    else
      b = c;
  }
  return S.compare(sa[b], T.length(), T) == 0 ? b : -1;
}
int contains2(const string &S, int *sa, const string &T) {
  int a = 0, b = S.length() + 1;
  while (b - a > 1) {
    int c = (a + b) / 2;
    if (S.compare(sa[c], T.length(), T) <= 0)
      a = c;
    else
      b = c;
  }
  return S.compare(sa[a], T.length(), T) == 0 ? a : -1;
}
int sa[MAX_N + 1], lcp[MAX_N + 1];

signed main() {
  string S, T, U;
  cin >> S >> T;
  int ans = 0;
  constract_sa(S, sa);

  for (int i = 0; i < (int)T.size(); i++) {
    char c = T[i];
    for (int j = 'a', k; j <= 'z'; j++) {
      if (j == c)
        continue;
      T[i] = j;
      if ((k = contains(S, sa, T)) >= 0)
        ans += contains2(S, sa, T) - k + 1;
    }
    for (int j = 'A', k; j <= 'Z'; j++) {
      if (j == c)
        continue;
      T[i] = j;
      if ((k = contains(S, sa, T)) >= 0)
        ans += contains2(S, sa, T) - k + 1;
    }
    T[i] = c;
    U += c;
    if (contains(S, sa, U) < 0)
      break;
  }
  cout << ans << endl;
  return 0;
}