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

/* #region header */
#ifdef LOCAL
#include "/Users/takakurashokichi/Desktop/atcoder/cxx-prettyprint-master/prettyprint.hpp"
#define debug(x) cout << x << endl
#else
#define debug(...) 42
#endif

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
// types
using namespace std;
using ll = long long;
using ull = unsigned long long;
using ld = long double;
typedef pair<ll, ll> Pl;
typedef pair<int, int> Pi;
typedef vector<ll> vl;
typedef vector<int> vi;
typedef vector<char> vc;
template <typename T> using mat = vector<vector<T>>;
typedef vector<vector<int>> vvi;
typedef vector<vector<long long>> vvl;
typedef vector<vector<char>> vvc;

template <int mod> struct modint {
  int x;

  modint() : x(0) {}

  modint(int64_t y) : x(y >= 0 ? y % mod : (mod - (-y) % mod) % mod) {}

  modint &operator+=(const modint &p) {
    if ((x += p.x) >= mod)
      x -= mod;
    return *this;
  }

  modint &operator-=(const modint &p) {
    if ((x += mod - p.x) >= mod)
      x -= mod;
    return *this;
  }

  modint &operator*=(const modint &p) {
    x = (int)(1LL * x * p.x % mod);
    return *this;
  }

  modint &operator/=(const modint &p) {
    *this *= p.inverse();
    return *this;
  }

  modint operator-() const { return modint(-x); }

  modint operator+(const modint &p) const { return modint(*this) += p; }

  modint operator-(const modint &p) const { return modint(*this) -= p; }

  modint operator*(const modint &p) const { return modint(*this) *= p; }

  modint operator/(const modint &p) const { return modint(*this) /= p; }

  bool operator==(const modint &p) const { return x == p.x; }

  bool operator!=(const modint &p) const { return x != p.x; }

  modint inverse() const {
    int a = x, b = mod, u = 1, v = 0, t;
    while (b > 0) {
      t = a / b;
      swap(a -= t * b, b);
      swap(u -= t * v, v);
    }
    return modint(u);
  }

  modint pow(int64_t n) const {
    modint ret(1), mul(x);
    while (n > 0) {
      if (n & 1)
        ret *= mul;
      mul *= mul;
      n >>= 1;
    }
    return ret;
  }

  friend ostream &operator<<(ostream &os, const modint &p) { return os << p.x; }

  friend istream &operator>>(istream &is, modint &a) {
    int64_t t;
    is >> t;
    a = modint<mod>(t);
    return (is);
  }

  static int get_mod() { return mod; }
};
// abreviations
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
#define rep_(i, a_, b_, a, b, ...) for (ll i = (a), max_i = (b); i < max_i; i++)
#define rep(i, ...) rep_(i, __VA_ARGS__, __VA_ARGS__, 0, __VA_ARGS__)
#define rrep_(i, a_, b_, a, b, ...)                                            \\
  for (ll i = (b - 1), min_i = (a); i >= min_i; i--)
#define rrep(i, ...) rrep_(i, __VA_ARGS__, __VA_ARGS__, 0, __VA_ARGS__)
#define SZ(x) ((ll)(x).size())
#define pb(x) push_back(x)
#define eb(x) emplace_back(x)
#define mp make_pair
#define print(x) cout << x << endl
#define vprint(x)                                                              \\
  rep(i, x.size()) cout << x[i] << ' ';                                        \\
  cout << endl
#define vsum(x) accumulate(all(x), 0LL)
#define vmax(a) *max_element(all(a))
#define vmin(a) *min_element(all(a))
#define lb(c, x) distance((c).begin(), lower_bound(all(c), (x)))
#define ub(c, x) distance((c).begin(), upper_bound(all(c), (x)))
// functions
// gcd(0, x) fails.
ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }
ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> T mypow(T x, ll n) {
  T ret = 1;
  while (n > 0) {
    if (n & 1)
      (ret *= x);
    (x *= x);
    n >>= 1;
  }
  return ret;
}
ll modpow(ll x, ll n, const ll mod) {
  ll ret = 1;
  while (n > 0) {
    if (n & 1)
      (ret *= x);
    (x *= x);
    n >>= 1;
    x %= mod;
    ret %= mod;
  }
  return ret;
}
uint64_t my_rand(void) {
  static uint64_t x = 88172645463325252ULL;
  x = x ^ (x << 13);
  x = x ^ (x >> 7);
  return x = x ^ (x << 17);
}
ll popcnt(ull x) { return __builtin_popcountll(x); }
// graph template
template <typename T> struct edge {
  int src, to;
  T cost;

  edge(int to, T cost) : src(-1), to(to), cost(cost) {}

  edge(int src, int to, T cost) : src(src), to(to), cost(cost) {}

  edge &operator=(const int &x) {
    to = x;
    return *this;
  }

  bool operator<(const edge<T> &r) const { return cost < r.cost; }

  operator int() const { return to; }
};
template <typename T> using Edges = vector<edge<T>>;
template <typename T> using WeightedGraph = vector<Edges<T>>;
using UnWeightedGraph = vector<vector<int>>;
struct Timer {
  clock_t start_time;
  void start() { start_time = clock(); }
  int lap() {
    // return x ms.
    return (clock() - start_time) * 1000 / CLOCKS_PER_SEC;
  }
};
/* #endregion*/
// constant
#define inf 1000000005
#define INF 4000000004000000000LL
#define mod 1000000007LL
#define endl '\\n'
typedef modint<mod> mint;
const long double eps = 0.000001;
const long double PI = 3.141592653589793;
// library
template <typename Monoid> struct SegmentTree {
  using F = function<Monoid(Monoid, Monoid)>;

  int sz;
  vector<Monoid> seg;

  const F f;
  const Monoid M1;

  SegmentTree(int n, const F f, const Monoid &M1) : f(f), M1(M1) {
    sz = 1;
    while (sz < n)
      sz <<= 1;
    seg.assign(2 * sz, M1);
  }

  void set(int k, const Monoid &x) { seg[k + sz] = x; }

  void build() {
    for (int k = sz - 1; k > 0; k--) {
      seg[k] = f(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  void update(int k, const Monoid &x) {
    k += sz;
    seg[k] = x;
    while (k >>= 1) {
      seg[k] = f(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  Monoid query(int a, int b) {
    Monoid L = M1, R = M1;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1)
        L = f(L, seg[a++]);
      if (b & 1)
        R = f(seg[--b], R);
    }
    return f(L, R);
  }

  Monoid operator[](const int &k) const { return seg[k + sz]; }

  template <typename C>
  int find_subtree(int a, const C &check, Monoid &M, bool type) {
    while (a < sz) {
      Monoid nxt = type ? f(seg[2 * a + type], M) : f(M, seg[2 * a + type]);
      if (check(nxt))
        a = 2 * a + type;
      else
        M = nxt, a = 2 * a + 1 - type;
    }
    return a - sz;
  }

  // check(seg[i])を満たす最小のb<=iを返す.なければ-1
  template <typename C> int find_first(int a, const C &check) {
    Monoid L = M1;
    if (a <= 0) {
      if (check(f(L, seg[1])))
        return find_subtree(1, check, L, false);
      return -1;
    }
    int b = sz;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1) {
        Monoid nxt = f(L, seg[a]);
        if (check(nxt))
          return find_subtree(a, check, L, false);
        L = nxt;
        ++a;
      }
    }
    return -1;
  }

  // check(seg[i])を満たす最小のi<bを返す.なければ-1
  template <typename C> int find_last(int b, const C &check) {
    Monoid R = M1;
    if (b >= sz) {
      if (check(f(seg[1], R)))
        return find_subtree(1, check, R, true);
      return -1;
    }
    int a = sz;
    for (b += sz; a < b; a >>= 1, b >>= 1) {
      if (b & 1) {
        Monoid nxt = f(seg[--b], R);
        if (check(nxt))
          return find_subtree(b, check, R, true);
        R = nxt;
      }
    }
    return -1;
  }
};
int main() {
  cin.tie(0);
  ios::sync_with_stdio(0);
  cout << setprecision(20);
  ll n;
  cin >> n;
  vector<string> S(n);
  rep(i, n) {
    cin >> S[i];
    reverse(all(S[i]));
  }
  sort(all(S));
  vl id(n);
  iota(all(id), 0);
  sort(all(id), [&](int i, int j) { return S[i].size() > S[j].size(); });

  auto f = [&](ll x, ll y) { return x + y; };
  vector<SegmentTree<ll>> seg(26, SegmentTree<ll>(n, f, 0));

  unordered_set<ll> que;
  ll len = S[id[0]].size();
  ll ans = 0;
  for (int i : id) {
    if (len > S[i].size() - 1) {
      for (int j : que) {
        rep(k, S[i].size() - 1, len) seg[S[j][k] - 'a'].update(j, 1);
      }
    }
    chmin(len, (ll)S[i].size() - 1);
    que.insert(i);
    seg[S[i].back() - 'a'].update(i, 1);
    string x = "";
    rep(j, S[i].size() - 1) x += S[i][j];
    ll l = ub(S, x);
    ll r = ub(S, x + '{');
    ans += seg[S[i].back() - 'a'].query(l, r) - 1;
  }
  print(ans);
}


Error: TLE

Fixed submission:

/* #region header */
#ifdef LOCAL
#include "/Users/takakurashokichi/Desktop/atcoder/cxx-prettyprint-master/prettyprint.hpp"
#define debug(x) cout << x << endl
#else
#define debug(...) 42
#endif

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
// types
using namespace std;
using ll = long long;
using ull = unsigned long long;
using ld = long double;
typedef pair<ll, ll> Pl;
typedef pair<int, int> Pi;
typedef vector<ll> vl;
typedef vector<int> vi;
typedef vector<char> vc;
template <typename T> using mat = vector<vector<T>>;
typedef vector<vector<int>> vvi;
typedef vector<vector<long long>> vvl;
typedef vector<vector<char>> vvc;

template <int mod> struct modint {
  int x;

  modint() : x(0) {}

  modint(int64_t y) : x(y >= 0 ? y % mod : (mod - (-y) % mod) % mod) {}

  modint &operator+=(const modint &p) {
    if ((x += p.x) >= mod)
      x -= mod;
    return *this;
  }

  modint &operator-=(const modint &p) {
    if ((x += mod - p.x) >= mod)
      x -= mod;
    return *this;
  }

  modint &operator*=(const modint &p) {
    x = (int)(1LL * x * p.x % mod);
    return *this;
  }

  modint &operator/=(const modint &p) {
    *this *= p.inverse();
    return *this;
  }

  modint operator-() const { return modint(-x); }

  modint operator+(const modint &p) const { return modint(*this) += p; }

  modint operator-(const modint &p) const { return modint(*this) -= p; }

  modint operator*(const modint &p) const { return modint(*this) *= p; }

  modint operator/(const modint &p) const { return modint(*this) /= p; }

  bool operator==(const modint &p) const { return x == p.x; }

  bool operator!=(const modint &p) const { return x != p.x; }

  modint inverse() const {
    int a = x, b = mod, u = 1, v = 0, t;
    while (b > 0) {
      t = a / b;
      swap(a -= t * b, b);
      swap(u -= t * v, v);
    }
    return modint(u);
  }

  modint pow(int64_t n) const {
    modint ret(1), mul(x);
    while (n > 0) {
      if (n & 1)
        ret *= mul;
      mul *= mul;
      n >>= 1;
    }
    return ret;
  }

  friend ostream &operator<<(ostream &os, const modint &p) { return os << p.x; }

  friend istream &operator>>(istream &is, modint &a) {
    int64_t t;
    is >> t;
    a = modint<mod>(t);
    return (is);
  }

  static int get_mod() { return mod; }
};
// abreviations
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
#define rep_(i, a_, b_, a, b, ...) for (ll i = (a), max_i = (b); i < max_i; i++)
#define rep(i, ...) rep_(i, __VA_ARGS__, __VA_ARGS__, 0, __VA_ARGS__)
#define rrep_(i, a_, b_, a, b, ...)                                            \\
  for (ll i = (b - 1), min_i = (a); i >= min_i; i--)
#define rrep(i, ...) rrep_(i, __VA_ARGS__, __VA_ARGS__, 0, __VA_ARGS__)
#define SZ(x) ((ll)(x).size())
#define pb(x) push_back(x)
#define eb(x) emplace_back(x)
#define mp make_pair
#define print(x) cout << x << endl
#define vprint(x)                                                              \\
  rep(i, x.size()) cout << x[i] << ' ';                                        \\
  cout << endl
#define vsum(x) accumulate(all(x), 0LL)
#define vmax(a) *max_element(all(a))
#define vmin(a) *min_element(all(a))
#define lb(c, x) distance((c).begin(), lower_bound(all(c), (x)))
#define ub(c, x) distance((c).begin(), upper_bound(all(c), (x)))
// functions
// gcd(0, x) fails.
ll gcd(ll a, ll b) { return b ? gcd(b, a % b) : a; }
ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> T mypow(T x, ll n) {
  T ret = 1;
  while (n > 0) {
    if (n & 1)
      (ret *= x);
    (x *= x);
    n >>= 1;
  }
  return ret;
}
ll modpow(ll x, ll n, const ll mod) {
  ll ret = 1;
  while (n > 0) {
    if (n & 1)
      (ret *= x);
    (x *= x);
    n >>= 1;
    x %= mod;
    ret %= mod;
  }
  return ret;
}
uint64_t my_rand(void) {
  static uint64_t x = 88172645463325252ULL;
  x = x ^ (x << 13);
  x = x ^ (x >> 7);
  return x = x ^ (x << 17);
}
ll popcnt(ull x) { return __builtin_popcountll(x); }
// graph template
template <typename T> struct edge {
  int src, to;
  T cost;

  edge(int to, T cost) : src(-1), to(to), cost(cost) {}

  edge(int src, int to, T cost) : src(src), to(to), cost(cost) {}

  edge &operator=(const int &x) {
    to = x;
    return *this;
  }

  bool operator<(const edge<T> &r) const { return cost < r.cost; }

  operator int() const { return to; }
};
template <typename T> using Edges = vector<edge<T>>;
template <typename T> using WeightedGraph = vector<Edges<T>>;
using UnWeightedGraph = vector<vector<int>>;
struct Timer {
  clock_t start_time;
  void start() { start_time = clock(); }
  int lap() {
    // return x ms.
    return (clock() - start_time) * 1000 / CLOCKS_PER_SEC;
  }
};
/* #endregion*/
// constant
#define inf 1000000005
#define INF 4000000004000000000LL
#define mod 1000000007LL
#define endl '\\n'
typedef modint<mod> mint;
const long double eps = 0.000001;
const long double PI = 3.141592653589793;
// library
template <typename Monoid> struct SegmentTree {
  using F = function<Monoid(Monoid, Monoid)>;

  int sz;
  vector<Monoid> seg;

  const F f;
  const Monoid M1;

  SegmentTree(int n, const F f, const Monoid &M1) : f(f), M1(M1) {
    sz = 1;
    while (sz < n)
      sz <<= 1;
    seg.assign(2 * sz, M1);
  }

  void set(int k, const Monoid &x) { seg[k + sz] = x; }

  void build() {
    for (int k = sz - 1; k > 0; k--) {
      seg[k] = f(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  void update(int k, const Monoid &x) {
    k += sz;
    seg[k] = x;
    while (k >>= 1) {
      seg[k] = f(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  Monoid query(int a, int b) {
    Monoid L = M1, R = M1;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1)
        L = f(L, seg[a++]);
      if (b & 1)
        R = f(seg[--b], R);
    }
    return f(L, R);
  }

  Monoid operator[](const int &k) const { return seg[k + sz]; }

  template <typename C>
  int find_subtree(int a, const C &check, Monoid &M, bool type) {
    while (a < sz) {
      Monoid nxt = type ? f(seg[2 * a + type], M) : f(M, seg[2 * a + type]);
      if (check(nxt))
        a = 2 * a + type;
      else
        M = nxt, a = 2 * a + 1 - type;
    }
    return a - sz;
  }

  // check(seg[i])を満たす最小のb<=iを返す.なければ-1
  template <typename C> int find_first(int a, const C &check) {
    Monoid L = M1;
    if (a <= 0) {
      if (check(f(L, seg[1])))
        return find_subtree(1, check, L, false);
      return -1;
    }
    int b = sz;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1) {
        Monoid nxt = f(L, seg[a]);
        if (check(nxt))
          return find_subtree(a, check, L, false);
        L = nxt;
        ++a;
      }
    }
    return -1;
  }

  // check(seg[i])を満たす最小のi<bを返す.なければ-1
  template <typename C> int find_last(int b, const C &check) {
    Monoid R = M1;
    if (b >= sz) {
      if (check(f(seg[1], R)))
        return find_subtree(1, check, R, true);
      return -1;
    }
    int a = sz;
    for (b += sz; a < b; a >>= 1, b >>= 1) {
      if (b & 1) {
        Monoid nxt = f(seg[--b], R);
        if (check(nxt))
          return find_subtree(b, check, R, true);
        R = nxt;
      }
    }
    return -1;
  }
};
int main() {
  cin.tie(0);
  ios::sync_with_stdio(0);
  cout << setprecision(20);
  ll n;
  cin >> n;
  vector<string> S(n);
  rep(i, n) {
    cin >> S[i];
    reverse(all(S[i]));
  }
  sort(all(S));
  vl id(n);
  iota(all(id), 0);
  sort(all(id), [&](int i, int j) { return S[i].size() > S[j].size(); });

  auto f = [&](ll x, ll y) { return x + y; };
  vector<SegmentTree<ll>> seg(26, SegmentTree<ll>(n, f, 0));

  unordered_set<ll> que;
  ll len = S[id[0]].size();
  ll ans = 0;
  for (int i : id) {
    for (int j : que) {
      rep(k, S[i].size() - 1, len) seg[S[j][k] - 'a'].update(j, 1);
    }
    chmin(len, (ll)S[i].size() - 1);
    que.insert(i);
    seg[S[i].back() - 'a'].update(i, 1);
    string x = "";
    rep(j, S[i].size() - 1) x += S[i][j];
    ll l = ub(S, x);
    ll r = ub(S, x + '{');
    ans += seg[S[i].back() - 'a'].query(l, r) - 1;
  }
  print(ans);
}