ABC213F - Common Prefixes

問題

https://atcoder.jp/contests/abc213/tasks/abc213_f

略説

Sの接尾辞配列で考える(あとで帳尻を合わせる)

sa := suffix_array(S) (O(N)で計算可能)

lcp := f(sa[i], sa[i+1]) (O(N)で計算可能)

f(l,r) = min(lcp[l, r))

g(l, r) := min(lcp[l, r)) (表記の簡略化のため定義した)

ans[i] = {g(0,i) + g(1,i) + ...} + g(i, i) + {... + g(i, N-2) + g(i, N-1)}

これはstackなど用いてans[0]~ans[N-1]まで合計O(N)で計算可能

stackの中を常に降順に保ちつつlcp[i]を投げていく

計算量はO(N)となる

ACコード

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

#include <atcoder/string>

int N;
string S;
long ans[1<<20];

int main() {
    cin >> N >> S;
    vector<int> sa = atcoder::suffix_array(S);
    vector<int> lcp = atcoder::lcp_array(S, sa);
    stack<pair<int,int>> stk;
    long sum = 0;
    for (int i=0; i<N-1; i++) {
        int cnt = 1;
        while (!stk.empty()) {
            auto [t, c] = stk.top();
            if (t < lcp[i]) break;
            stk.pop();
            sum -= 1L*t*c;
            cnt += c;
        }
        stk.push(make_pair(lcp[i], cnt));
        sum += 1L*lcp[i]*cnt;
        ans[sa[i+1]] += sum;
    }
    while (!stk.empty()) stk.pop();
    sum = 0;
    for (int i=N-2; i>=0; i--) {
        int cnt = 1;
        while (!stk.empty()) {
            auto [t, c] = stk.top();
            if (t < lcp[i]) break;
            stk.pop();
            sum -= 1L*t*c;
            cnt += c;
        }
        stk.push(make_pair(lcp[i], cnt));
        sum += 1L*lcp[i]*cnt;
        ans[sa[i]] += sum;
    }
    for (int i=0; i<N; i++) cout << ans[i]+N-i << '\n';
}