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

#include <random>

#include "bits/stdc++.h"

using namespace std;

using ll = long long;
using ld = long double;
using P = pair<int, int>;

template <typename T> void printv(const vector<T> &v) {
  int sz = v.size();
  for (int i = 0; i < sz; i++) {
    cout << v[i] << " \\n"[i == sz - 1];
  }
}

struct Trie {
  Trie *node[26];
  vector<ll> val;
  Trie() : val(26, 0) { std::fill(node, node + 26, nullptr); }
  void insert(const std::string &s) {
    Trie *r = this;
    int sz = s.size();
    for (int i = 0; i < sz - 1; i++) {
      int c = s[i] - 'a';
      if (!r->node[c])
        r->node[c] = new Trie;
      r = r->node[c];
    }
    int idx = s[sz - 1] - 'a';
    r->val[idx]++;
    return;
  }
  Trie *find(const char &s, Trie *pos) {
    int c = s - 'a';
    if (!pos->node[c])
      return nullptr;
    return pos->node[c];
  }
  ll count(const string &s) {
    ll ans = 0;
    vector<int> cnt(26);
    for (auto c : s)
      cnt[c - 'a']++;
    Trie *r = this;
    for (int i = 0; i < (int)s.size(); i++) {
      for (int j = 0; j < 26; j++) {
        if (cnt[j]) {
          ans += r->val[j];
        }
      }
      auto nxt = find(s[i], r);
      if (!nxt)
        break;
      r = nxt;
      cnt[s[i] - 'a']--;
    }
    return ans;
  }
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  cin >> n;
  vector<string> s(n);
  for (int i = 0; i < n; i++) {
    cin >> s[i];
    reverse(s[i].begin(), s[i].end());
  }
  sort(s.begin(), s.end(),
       [](const string &x, const string &y) { return x.size() < y.size(); });
  Trie trie;
  ll ans = 0;
  for (int i = 0; i < n; i++) {
    ans += trie.count(s[i]);
    trie.insert(s[i]);
  }
  cout << ans << endl;
}


Error: TLE

Fixed submission:

#include <random>

// C++ includes used for precompiling -*- C++ -*-

// Copyright (C) 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010
// Free Software Foundation, Inc.
//
// This file is part of the GNU ISO C++ Library.  This library is free
// software; you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the
// Free Software Foundation; either version 3, or (at your option)
// any later version.

// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// Under Section 7 of GPL version 3, you are granted additional
// permissions described in the GCC Runtime Library Exception, version
// 3.1, as published by the Free Software Foundation.

// You should have received a copy of the GNU General Public License and
// a copy of the GCC Runtime Library Exception along with this program;
// see the files COPYING3 and COPYING.RUNTIME respectively.  If not, see
// <http://www.gnu.org/licenses/>.

/** @file stdc++.h
 *  This is an implementation file for a precompiled header.
 */

// 17.4.1.2 Headers

// C
#ifndef _GLIBCXX_NO_ASSERT
#include <cassert>
#endif
#include <cctype>
#include <cerrno>
#include <cfloat>
#include <ciso646>
#include <climits>
#include <clocale>
#include <cmath>
#include <csetjmp>
#include <csignal>
#include <cstdarg>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>

#ifdef __GXX_EXPERIMENTAL_CXX0X__
#include <ccomplex>
#include <cfenv>
#include <cinttypes>
#include <cstdbool>
#include <cstdint>
#include <ctgmath>
#include <cwchar>
#include <cwctype>
#endif

// C++
#include <algorithm>
#include <bitset>
#include <complex>
#include <deque>
#include <exception>
#include <fstream>
#include <functional>
#include <iomanip>
#include <ios>
#include <iosfwd>
#include <iostream>
#include <istream>
#include <iterator>
#include <limits>
#include <list>
#include <locale>
#include <map>
#include <memory>
#include <new>
#include <numeric>
#include <ostream>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdexcept>
#include <streambuf>
#include <string>
#include <typeinfo>
#include <utility>
#include <valarray>
#include <vector>

#ifdef __GXX_EXPERIMENTAL_CXX0X__
#include <array>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <forward_list>
#include <future>
#include <initializer_list>
#include <mutex>
#include <random>
#include <ratio>
#include <regex>
#include <system_error>
#include <thread>
#include <tuple>
#include <type_traits>
#include <typeindex>
#include <unordered_map>
#include <unordered_set>
#endif

using namespace std;

using ll = long long;
using ld = long double;
using P = pair<uint64_t, uint64_t>;

template <typename T> void printv(const vector<T> &v) {
  int sz = v.size();
  for (int i = 0; i < sz; i++) {
    cout << v[i] << " \\n"[i == sz - 1];
  }
}

struct XorShift {
  using result_type = uint32_t;
  result_type w = 123456789, x = 362436069, y = 521288629, z = 88675123;
  XorShift(result_type seed = time(nullptr)) {
    w = seed;
    x = w << 13;
    y = (w >> 9) ^ (x << 6);
    z = y >> 7;
  }
  static constexpr result_type min() { return 0; }
  static constexpr result_type max() { return 0x7FFFFFFF; }
  result_type operator()() {
    result_type t = x ^ (x << 11);
    x = y;
    y = z;
    z = w;
    return w = (w ^ (w >> 19) ^ (t ^ (t >> 8)));
  }
  result_type rand() {
    result_type t = x ^ (x << 11);
    x = y;
    y = z;
    z = w;
    return w = (w ^ (w >> 19) ^ (t ^ (t >> 8)));
  }
  // [min,max] の整数値乱数
  result_type randInt(result_type min = 0, result_type max = 0x7FFFFFFF) {
    return rand() % (max - min + 1) + min;
  }
  // [min,max] の浮動小数点乱数
  double randDouble(double min = 0, double max = 1) {
    return (double)(rand() % 0xFFFF) / 0xFFFF * (max - min) + min;
  }
};

XorShift rnd(114514);

uint64_t b1 = (uint64_t)rnd.randInt(2, 114514),
         b2 = (uint64_t)rnd.randInt(2, 114514);

class RollingHash {
  int n;
  static const uint64_t moda = (uint64_t)1e9 + 7;
  static const uint64_t modb = (uint64_t)1e9 + 9;
  // static const uint64_t b1 = (uint64_t)17, b2 = (uint64_t)19;
  std::vector<uint64_t> a, b, p, q;

public:
  RollingHash(const std::string &str)
      : n((int)str.size()), a(n + 1, 0), b(n + 1, 0), p(n + 1, 1), q(n + 1, 1) {
    for (int i = 0; i < n; i++) {
      a[i + 1] = (a[i] * b1 + uint64_t(str[i])) % moda;
      b[i + 1] = (b[i] * b2 + uint64_t(str[i])) % modb;
      p[i + 1] = p[i] * b1 % moda;
      q[i + 1] = q[i] * b2 % modb;
    }
  }
  // [l, r) のハッシュ値を取得
  std::pair<uint64_t, uint64_t> query(int l, int r) {
    assert(0 <= l && l <= r && r <= n);
    uint64_t h1 = (p[r - l] * a[l] - a[r] + moda) % moda;
    uint64_t h2 = (q[r - l] * b[l] - b[r] + modb) % modb;
    return std::make_pair(h1, h2);
  }
};

map<P, ll> c_sufs[26];
ll c_cnt[26];

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  cin >> n;
  vector<string> s(n);
  for (int i = 0; i < n; i++) {
    cin >> s[i];
  }
  sort(s.begin(), s.end(),
       [](const string &x, const string &y) { return x.size() > y.size(); });

  ll ans = 0;
  for (int i = 0; i < n; i++) {
    // cerr << s[i] << endl;
    if (s[i].size() == 1) {
      int idx = s[i][0] - 'a';
      ans += c_cnt[idx];
      c_cnt[idx]++;
      continue;
    }
    vector<int> ok(26);
    for (auto c : s[i]) {
      int idx = c - 'a';
      ok[idx] = 1;
    }
    for (int j = 0; j < 26; j++)
      if (ok[j])
        c_cnt[j]++;
    RollingHash rh(s[i]);
    int sz = s[i].size();

    // [][----]
    auto p = rh.query(1, sz);
    int index = s[i][0] - 'a';
    if (c_sufs[index].count(p))
      ans += c_sufs[index][p];

    vector<int> cnt(26);
    for (int j = 0; j < sz; j++) {
      int idx = s[i][j] - 'a';
      auto val = rh.query(j, sz);
      // cerr << j << " " << val.first << " " << val.second << endl;
      for (int k = 0; k < 26; k++) {
        if (cnt[k])
          c_sufs[k][val]++;
      }
      cnt[idx]++;
    }
  }
  cout << ans << endl;
}