はまやんはまやんはまやん

hamayanhamayan's blog

Common Prefixes [AtCoder Beginner Contest 213 F]

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

前提知識

解説

https://atcoder.jp/contests/abc213/submissions/24898387

Suffix ArrayとLCPの知識が無ければ、まず解くことはできない。
以下で軽く説明するが、他方で勉強する方がオススメ。

愚直解

問題で要求されている問題をそのまま解くのも意外と難しく、前提知識が要求される。
なので、まずは愚直解が解けるかどうかがポイントになる。
f(S[k], S[i])を高速に解くための方法がSuffix Array+LCPになる。

前提知識1: Suffix Array

これは文字列Sのsuffix(接尾語)を辞書順に並べた配列のことである。
接尾語?という感じかもしれないが、今回の問題は接尾語の集合に対する操作になっている。
abcにはabc, bc, cという接尾語のレパートリーがあり、それらの類似度を独自に計算している。

Suffix Arrayの雰囲気を軽く説明する。
aabaacという文字列に対してSuffix Arrayを考えてみる。
接尾語を書き下すと

aabaac  
abaac  
baac  
aac  
ac  
c  

となる。これを辞書順に並べると

aabaac  
aac  
abaac  
ac  
baac  
c  

となる。これを得るのが、Suffix Array
本当は、具体的にはどこから始まったかだけが得られるので{0,3,1,4,2,1}が得られる。

前提知識2: LCP

LCPは今回の類似度で要求されている先頭から何文字同じかを求めるものであり、Suffix Arrayの場合であれば、
辞書順に並べたときに隣接する文字列で何文字同じかを求めたものになる。

aabaac  
aac  
abaac  
ac  
baac  
c  

この例を再度使うと、隣接する文字列で何文字同じかを見ればいいので

2  
1  
1  
0  
0  

となる。
このSuffix ArrayとLCPと区間最小値のセグメントツリー(またはスパーステーブル)によって任意の接頭語同士の
先頭から同じ文字数を高速に求めることができる。
例えば、aabaacとabaacの同じ文字数を求めたいときは、まずはSuffix Arrayのどことどこにあるかを見つけて

> aabaac  
  aac  
> abaac  
  ac  
  baac  
  c  

その間にあるLCPの最小値を求めれば、{2, 1}なので1であり、それが2つの接尾語の共通する先頭文字列数となる。
このようにSuffix ArrayとLCPが求められていれば、2つの接尾語の共通する先頭文字列数、f(S[k], S[i])は
区間minが計算できれば求めることができるようになる。
区間minをスパーステーブルを使うことができれば、愚直解はO(N2)で計算することができる。
ここまでをしっかり理解できる所が折り返し地点となる。

なお、Suffix ArrayとLCPについてはAtCoder Libraryにライブラリがあるので活用するといい。
区間minの部分については自分で何とかする。
自分は丸ごとライブラリにしてある。

高速化へ

高速化への糸口を探すために以下の例を使おう。
Suffix Arrayの順番になっているので、答えるときはもともとの順番に変換する必要があるが、そこは頑張って。

1 aabaac  
2 aac  
3 abaac  
4 ac  
5 baac  
6 c  

f(S[k], S[i])を計算していくが、ここでi<kに限定して計算していくことをまずは考えていく。
i=kについては事前に文字数分足し合わせておくことにして、i>kについてはこれから説明することを順番をひっくり返してやればいい。
まあ、とにかくi<kの場合に限定して考えてみる。

k=1の場合は特に何もしない。

k=2の場合はf(S[2], S[1])を計算すればいい。これはこれまでに説明したことからできる。

k=3の場合はf(S[3], S[2])+f(S[3], S[1])を計算すればいい。
この後半部分は、最小値が計算されているということを考慮すると、

f(S[3], S[1]) = min(f(S[3], S[2]), f(S[2], S[1]))

となる。これがすごい大事なので、ちゃんと理解してほしい。
f(S[3], S[1])はLCP的にはmin(lcp[3][2], lcp[2][1])を取っているような感じになるので、
f(S[3],S[2])=lcp[3][2], f(S[2],S[1])=lcp[2][1]であることを考えると筋が通る。

k=4の場合はf(S[4], S[3])+f(S[4], S[2])+f(S[4], S[1])を計算すればいい。
これも同様に考えてみると、
f(S[4], S[2]) = min( f(S[4],S[3]), f(S[3], S[2]) )
f(S[4], S[1]) = min( f(S[4],S[3]), f(S[3], S[1]) )
となる。

さて、これがなんだ?という話に戻るが、とあるkを処理する場合にk-1から高速に差分計算をすることで
f(S[k], S[k-1]) + f(S[k], S[k-2]) + ... + f(S[k], S[1])を計算することができるということである。

f(S[k-1], S[k-2]) + f(S[k-1], S[k-3]) + ... + f(S[k-1], S[1])

がもともとあったとして、すべての要素に対してf(S[k], S[k-1])とminを取る。すると、

f(S[k], S[k-2]) + f(S[k], S[k-3]) + ... + f(S[k], S[1])

こういう風に値が変化するので、あとは、f(S[k], S[k-1])を足せば

f(S[k], S[k-1]) + f(S[k], S[k-2]) + f(S[k], S[k-3]) + ... + f(S[k], S[1])

となって差分計算ができる。この理解が最後の山になるので頑張ってほしい。
これが分かれば、あとは、

  • 全体にmin xを取る
  • xを入れる
  • 総和を求める

ことができるデータ構造があれば答えを求めることができる。
ここまで理解できれば、もうほぼ解けているし、ここまで分かる人にとってはデータ構造づくりは大丈夫。

データ構造

自分はstackを使って実装した。
{数, その数の個数}をstackに入れる感じで実装した。
別途、stackに入っている数の総和をtotとしてもっておく。
stackには深いほど小さい数が入っているような感じにしておいて、新しい数xを入れるときに、
xより大きい数は取り出してtotから引いて、その数をxに変換してstackとtotに入れなおすということをする。
つまり、stackには狭義単調減少している状態を保ちながら、総和を持っておくシステムである。

こうすることで毎回min操作でstackに入れる回数は1回になり、stackから出す回数は入れる回数を上回るわけはないので、
stackに入れる回数がO(N)で抑えられるので全体としてはO(N)で実装が行える。

int N; string S;
ll ans[1010101];
//---------------------------------------------------------------------------------------------------
stack<pair<int, int>> que;
ll tot = 0;
void init() {
    while (!que.empty()) que.pop();
    tot = 0;
}
void add(int x) {
    int cnt = 1;
    while (!que.empty()) {
        if (que.top().first < x) break;

        auto q = que.top(); que.pop();
        tot -= 1LL * q.first * q.second;
        cnt += q.second;
    }

    tot += 1LL * x * cnt;
    que.push({ x, cnt });
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> S;

    SuffixArraySAIS sais(S);

    rep(i, 1, N + 1) ans[sais.sa[i]] = N - sais.sa[i];
    rep(i, 2, N + 1) {
        add(sais.common_prefix(sais.sa[i - 1], sais.sa[i]));
        ans[sais.sa[i]] += tot;
    }
    
    init();
    rrep(i, N - 1, 1) {
        add(sais.common_prefix(sais.sa[i + 1], sais.sa[i]));
        ans[sais.sa[i]] += tot;
    }

    rep(i, 0, N) printf("%lld\n", ans[i]);
}