Submission #3369820

---

Source Code Expand

#include <bits/stdc++.h>
#define ll long long
#define REP(i, n) for (ll (i) = 0; (i) < (n); (i)++)
#define REPI(i, a, b) for (ll (i) = (a); (i) < (b); (i)++)
#define int long long
using namespace std;
using II = pair<int, int>;
using VI = vector<int>;
using VVI = vector<VI>;
using VVVI = vector<VVI>;
using VII = vector<II>;

class SuffixArray {
  void set_lcp() {
    lcp = VI(N - 1);
    int h = 0;
    REP (i, N) {
      if (h > 0) h--;
      if (rank[i] == N - 1) continue;
      int j = sorted[rank[i] + 1]; // 比べる対象(辞書順が一つ大きいもの)
      for (; i + h < N && j + h < N; h++) {
        if (S[i + h] != S[j + h]) break;
      }
      lcp[rank[i]] = h;
    }
  }
public:
  int N;
  string S;
  VI rank; // rank[i]: iから始まるsuffixの辞書順での順位
  VI sorted; // sorted[i]: suffixが辞書順i番目となる開始位置のindex
  VI lcp; // lcp[i]: S[sorted[i]...]とS[sorted[i + 1]...]が先頭何文字一致しているか
  SuffixArray(string s) {
    S = s;
    N = S.size();
    sorted = VI(N); rank = VI(N);
    REP (i, N) {
      sorted[i] = i;
      rank[i] = S[i];
    }

    int k;
    function<bool(int, int)> compare_sa = [this, &k](int i, int j) {
      if (rank[i] != rank[j]) { return rank[i] < rank[j]; }
      int ri = i + k < N ? rank[i + k] : -1;
      int rj = j + k < N ? rank[j + k] : -1;
      return ri < rj;
    };

    for (k = 1; k < N; k *= 2) {
      sort(sorted.begin(), sorted.end(), compare_sa);
      VI tmp(N, 0);
      REPI (i, 1, N) {
        tmp[sorted[i]] = tmp[sorted[i - 1]] + compare_sa(sorted[i - 1], sorted[i]);
      }
      rank = tmp;
    }
    set_lcp();
  }

  // sizeがTと等しく初めてT以上になるようなSの部分文字列(sortedのイテレータを返す)
  VI::iterator lower_bound(string T) {
    int l = -1, r = N;
    while (r - l > 1) {
      int mid = (l + r) / 2;
      if (S.compare(sorted[mid], T.size(), T) < 0) {
        l = mid;
      } else {
        r = mid;
      }
    }
    return sorted.begin() + r;
  }

  // sizeがTと等しく初めてTより大きくなるようなSの部分文字列(sortedのイテレータを返す)
  VI::iterator upper_bound(string T) {
    int l = -1, r = N;
    while (r - l > 1) {
      int mid = (l + r) / 2;
      if (S.compare(sorted[mid], T.size(), T) <= 0) {
        l = mid;
      } else {
        r = mid;
      }
    }
    return sorted.begin() + r;
  }

  // S2が部分文字列として何回出現するか
  int count(string S2) {
    return upper_bound(S2) - lower_bound(S2);
  }
};

signed main() {
  string S; cin >> S;
  SuffixArray sa(S);
  int N = S.size();
  int ans = N * (N + 1) * (N + 2) / 6;
  REP (i, N - 1) {
    ans -= sa.lcp[i] * (sa.lcp[i] + 1) / 2;
  }
  cout << ans << endl;
}

Submission Info

Judge Result