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

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#define _USE_MATH_DEFINES
#include <complex>
#include <functional>
#include <math.h>
using namespace std;

#define rep(i, x) for (ll i = 0; i < x; i++)
#define repn(i, x) for (ll i = 1; i <= x; i++)

typedef long long ll;
const ll INF = 1e17;
const ll MOD = 1000000007;
const ll MAX = 4000001;
const long double eps = 1E-14;

ll max(ll a, ll b) {
  if (a > b) {
    return a;
  }
  return b;
}

ll min(ll a, ll b) {
  if (a > b) {
    return b;
  }
  return a;
}

ll gcd(ll a, ll b) {
  if (b == 0) {
    return a;
  }
  if (a < b) {
    return gcd(b, a);
  }
  return gcd(b, a % b);
}

ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }

struct edge {
  ll ind;
  ll fr;
  ll to;
  ll d;
};

class mint {
  long long x;

public:
  mint(long long x = 0) : x((x % MOD + MOD) % MOD) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint &a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator-=(const mint &a) {
    if ((x += MOD - a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator*=(const mint &a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  mint operator+(const mint &a) const {
    mint res(*this);
    return res += a;
  }
  mint operator-(const mint &a) const {
    mint res(*this);
    return res -= a;
  }
  mint operator*(const mint &a) const {
    mint res(*this);
    return res *= a;
  }

  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime MOD
  mint inv() const { return pow(MOD - 2); }
  mint &operator/=(const mint &a) { return (*this) *= a.inv(); }
  mint operator/(const mint &a) const {
    mint res(*this);
    return res /= a;
  }

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

mint pw(mint a, ll b) {
  if (b == 0) {
    return 1;
  }
  mint ret = pw(a, b >> 1);
  ret *= ret;
  if (b & 1) {
    ret *= a;
  }
  return ret;
}

typedef vector<ll> vll;
typedef vector<vector<ll>> vvll;
typedef vector<vector<vector<ll>>> vvvll;

typedef vector<mint> vmint;
typedef vector<vector<mint>> vvmint;
typedef vector<vector<vector<mint>>> vvvmint;

/////////////////////////////////////
struct TRIE {
  vector<string> S; // 配列は1-indexed
  ll N;

  vvll graph; // 0-indexed(0はroot)
  vector<char> ch;

  TRIE(vector<string> S0) {
    S = S0;
    ll N = S.size() - 1;

    repn(i, N) S[i] += '*';

    graph.push_back({});
    ch.push_back('.');

    repn(i, N) {
      ll v = 0;
      rep(j, S[i].size()) {
        ll newv = -1;

        for (ll w : graph[v]) {
          if (ch[w] == S[i][j]) {
            newv = w;
            break;
          }
        }

        if (newv == -1) {
          newv = graph.size();
          graph.push_back({});
          graph[v].push_back(newv);
          ch.push_back(S[i][j]);
        }
        v = newv;
      }
    }
  }
};

int main() {
  ll N;
  cin >> N;
  vector<string> S(N + 1);
  repn(i, N) {
    cin >> S[i];
    reverse(S[i].begin(), S[i].end());
  }

  TRIE t(S);
  vvll g = t.graph;
  vector<char> col = t.ch;
  ll M = g.size();
  // cout << M << endl;

  // rep(i, M)cout << col[i] << endl;
  // rep(v, M)for (ll w : g[v]) { cout << v << " "<<w << endl; }

  vll par(M, -1);
  vll sz(M, 0);
  for (ll v = M - 1; v >= 0; v--) {
    if (g[v].size() == 0) {
      sz[v] = 1;
    }
    for (ll w : g[v]) {
      par[w] = v;
      sz[v] += sz[w];
    }
  }

  vvll count(M, vll(26, 0));
  for (ll v = M - 1; v >= 0; v--) {
    for (ll w : g[v]) {
      rep(j, 26) count[v][j] += count[w][j];
    }

    if (col[v] != '.' && col[v] != '*') {
      count[v][col[v] - 'a'] = sz[v];
    }
  }

  ll ans = 0;
  rep(v, M) {
    if (col[v] != '*') {
      continue;
    }
    char c = col[par[v]];
    ll w = par[par[v]];

    // cout << w << " " << c << endl;
    for (ll u : g[w]) {
      ans += count[u][c - 'a'];
    }
    ans--;
    // ans += count[w][c - 'a'] - 1;
  }

  cout << ans << endl;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#define _USE_MATH_DEFINES
#include <complex>
#include <functional>
#include <math.h>
using namespace std;

#define rep(i, x) for (ll i = 0; i < x; i++)
#define repn(i, x) for (ll i = 1; i <= x; i++)

typedef long long ll;
const ll INF = 1e17;
const ll MOD = 1000000007;
const ll MAX = 4000001;
const long double eps = 1E-14;

ll max(ll a, ll b) {
  if (a > b) {
    return a;
  }
  return b;
}

ll min(ll a, ll b) {
  if (a > b) {
    return b;
  }
  return a;
}

ll gcd(ll a, ll b) {
  if (b == 0) {
    return a;
  }
  if (a < b) {
    return gcd(b, a);
  }
  return gcd(b, a % b);
}

ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }

struct edge {
  ll ind;
  ll fr;
  ll to;
  ll d;
};

class mint {
  long long x;

public:
  mint(long long x = 0) : x((x % MOD + MOD) % MOD) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint &a) {
    if ((x += a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator-=(const mint &a) {
    if ((x += MOD - a.x) >= MOD)
      x -= MOD;
    return *this;
  }
  mint &operator*=(const mint &a) {
    (x *= a.x) %= MOD;
    return *this;
  }
  mint operator+(const mint &a) const {
    mint res(*this);
    return res += a;
  }
  mint operator-(const mint &a) const {
    mint res(*this);
    return res -= a;
  }
  mint operator*(const mint &a) const {
    mint res(*this);
    return res *= a;
  }

  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime MOD
  mint inv() const { return pow(MOD - 2); }
  mint &operator/=(const mint &a) { return (*this) *= a.inv(); }
  mint operator/(const mint &a) const {
    mint res(*this);
    return res /= a;
  }

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

mint pw(mint a, ll b) {
  if (b == 0) {
    return 1;
  }
  mint ret = pw(a, b >> 1);
  ret *= ret;
  if (b & 1) {
    ret *= a;
  }
  return ret;
}

typedef vector<ll> vll;
typedef vector<vector<ll>> vvll;
typedef vector<vector<vector<ll>>> vvvll;

typedef vector<mint> vmint;
typedef vector<vector<mint>> vvmint;
typedef vector<vector<vector<mint>>> vvvmint;

/////////////////////////////////////

string S[200001];

int main() {
  ll N;
  cin >> N;

  vector<string> S(N + 1);
  map<string, int> mp;
  repn(i, N) {
    cin >> S[i];
    reverse(S[i].begin(), S[i].end());

    string T = S[i].substr(0, S[i].size() - 1);
    int c = S[i][S[i].size() - 1] - 'a';
    mp[T] += 1ull << c;

    // cout << T << mp[T] << endl;
  }

  ll ans = 0;
  repn(i, N) {

    int c = S[i][S[i].size() - 1] - 'a';
    int cnt = 1 << c;

    S[i].pop_back();

    while (1) {
      // cout << S[i] << endl;
      // cout << cnt << endl;
      // if (!mp.count(S[i])) {
      //	if (S[i].size() == 0) { break; }
      //	S[i].pop_back();

      //	continue;
      //}

      int x = mp[S[i]];

      // cout << x << endl;

      x &= cnt;
      while (x > 0) {
        if (x & 1) {
          ans++;
        }
        x = x >> 1;
      }

      if (S[i].size() == 0) {
        break;
      }

      int cc = S[i][S[i].size() - 1] - 'a';
      if (((cnt >> cc) & 1) == 0)
        cnt += 1ull << cc;

      S[i].pop_back();
    }
  }

  cout << ans - N << endl;
}