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

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

// #define int long long

typedef pair<int, int> pint;
typedef vector<int> vint;
typedef vector<pint> vpint;
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define all(v) (v).begin(), (v).end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, f, n) for (int i = (f); i < (n); i++)
#define each(it, v)                                                            \\
  for (__typeof((v).begin()) it = (v).begin(); it != (v).end(); it++)
template <class T, class U> void chmin(T &t, U f) {
  if (t > f)
    t = f;
}
template <class T, class U> void chmax(T &t, U f) {
  if (t < f)
    t = f;
}

namespace SA {
const int MAX = 1000000;
int N, K;
int tmp[MAX + 1], rank[MAX + 1];
bool compare_sa(int i, int j) {
  if (rank[i] != rank[j])
    return rank[i] < rank[j];
  else {
    int ri = (i + K <= N) ? rank[i + K] : -1;
    int rj = (j + K <= N) ? rank[j + K] : -1;
    return ri < rj;
  }
}

vector<int> construct_sa(string S) {
  N = S.size();
  vector<int> sa(N + 1);
  for (int i = 0; i <= N; i++)
    sa[i] = i;
  for (int i = 0; i <= N; i++) {
    rank[i] = i == N ? -1 : S[i];
  }

  for (K = 1; K <= N; K *= 2) {
    sort(sa.begin(), sa.end(), compare_sa);
    tmp[sa[0]] = 0;
    for (int i = 1; i <= N; i++) {
      tmp[sa[i]] = tmp[sa[i - 1]] + (compare_sa(sa[i - 1], sa[i]) ? 1 : 0);
    }
    for (int i = 0; i <= N; i++)
      rank[i] = tmp[i];
  }
  return sa;
}
} // namespace SA

bool exist(string &S, vint &sa, string &T) {
  int a = 0, b = S.size();
  while (b - a > 1) {
    int c = (a + b) / 2;
    if (S.substr(sa[c], S.size() - sa[c]) < T)
      a = c;
    else
      b = c;
  }
  return S.substr(sa[b], min(S.size() - sa[b], T.size())) == T;
}

bool contain(string &S, vint &sa, string &T) {
  int a = 0, b = S.size();
  while (b - a > 1) {
    int c = (a + b) / 2;
    if (S.compare(sa[c], T.size(), T) < 0)
      a = c;
    else
      b = c;
  }
  return S.compare(sa[b], T.size(), T) == 0;
}

signed main() {
  string S;
  vint sa;
  int Q;
  cin >> S;
  sa = SA::construct_sa(S);
  cin >> Q;

  while (Q--) {
    string T;
    cin >> T;
    if (contain(S, sa, T))
      cout << 1 << endl;
    else
      cout << 0 << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

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

// #define int long long

typedef pair<int, int> pint;
typedef vector<int> vint;
typedef vector<pint> vpint;
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define all(v) (v).begin(), (v).end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, f, n) for (int i = (f); i < (n); i++)
#define each(it, v)                                                            \\
  for (__typeof((v).begin()) it = (v).begin(); it != (v).end(); it++)
template <class T, class U> void chmin(T &t, U f) {
  if (t > f)
    t = f;
}
template <class T, class U> void chmax(T &t, U f) {
  if (t < f)
    t = f;
}

namespace SA {
const int MAX = 1000000;
int N, K;
int tmp[MAX + 1], rank[MAX + 1];
bool compare_sa(int i, int j) {
  if (rank[i] != rank[j])
    return rank[i] < rank[j];
  else {
    int ri = (i + K <= N) ? rank[i + K] : -1;
    int rj = (j + K <= N) ? rank[j + K] : -1;
    return ri < rj;
  }
}

vector<int> construct_sa(string S) {
  N = S.size();
  vector<int> sa(N + 1);
  for (int i = 0; i <= N; i++)
    sa[i] = i;
  for (int i = 0; i <= N; i++) {
    rank[i] = i == N ? -1 : S[i];
  }

  for (K = 1; K <= N; K *= 2) {
    sort(sa.begin(), sa.end(), compare_sa);
    tmp[sa[0]] = 0;
    for (int i = 1; i <= N; i++) {
      tmp[sa[i]] = tmp[sa[i - 1]] + (compare_sa(sa[i - 1], sa[i]) ? 1 : 0);
    }
    for (int i = 0; i <= N; i++)
      rank[i] = tmp[i];
  }
  return sa;
}
} // namespace SA

bool exist(string &S, vint &sa, string &T) {
  int a = 0, b = S.size();
  while (b - a > 1) {
    int c = (a + b) / 2;
    if (S.substr(sa[c], S.size() - sa[c]) < T)
      a = c;
    else
      b = c;
  }
  return S.substr(sa[b], min(S.size() - sa[b], T.size())) == T;
}

bool contain(string &S, vint &sa, string &T) {
  int a = 0, b = S.size();
  while (b - a > 1) {
    int c = (a + b) / 2;
    if (S.compare(sa[c], T.size(), T) < 0)
      a = c;
    else
      b = c;
  }
  return S.compare(sa[b], T.size(), T) == 0;
}

signed main() {
  string S;
  vint sa;
  int Q;
  cin >> S;
  sa = SA::construct_sa(S);
  cin >> Q;

  while (Q--) {
    string T;
    cin >> T;
    if (exist(S, sa, T))
      cout << 1 << endl;
    else
      cout << 0 << endl;
  }
  return 0;
}