以下では文字列 \(S\) の長さを \(|S|\) と表します。また、問題文中で定義したとおりに、類似度 という語、及び \(f\), \(S_i\) という記号を用います。

配列などの index は 0-based とします。

解法の方針

suffix たちの common prefix の長さの和を求める問題なので、LCP Array (Longest Common Prefix Array) を使うと解けそうです。

Suffix Array

文字列 \(S\) の suffix を辞書順に並べたものを考えます。

例えば aacab という文字列であれば、suffix たちは aacab, acab, cab, ab, b, (空文字列) であり、これを辞書順にソートすると (空文字列), aacab, ab, acab, b, cab となります。

こうしてソートされた配列を表すために、文字列をいちいち保持していると \(O(|S|^2)\) のメモリが必要になりますが、「\(S\) の何文字目以降からなる文字列か」だけを保持することにすれば、\(O(|S|)\) で済みます。上の例では \((5,0,3,1,4,2)\) となります。この数列のことを Suffix Array といいます。以下ではこの配列を \(SA\) と表します。

Suffix Array を得るために実際に文字列の大小比較をすると \(O(|S|^2\log |S|)\) の時間がかかってしまいますが、適切なアルゴリズムにより (文字の種類数を定数として) \(O(|S|)\) で構築することができます。（SA-IS と呼ばれるアルゴリズムです。詳細は省略します。また、\(O(|S|\log |S|)\) や \(O(|S|(\log |S|)^2)\) の時間が掛かるかわりに実装が比較的容易なアルゴリズムもあります）

LCP Array

辞書順で隣接するsuffixの類似度を考えます。例えば前述の例であれば

• (空文字列) と aacab の類似度は \(0\)
• aacab とab の類似度は \(1\)
• ab と acab の類似度は \(1\)
• acab と b の類似度は \(0\)
• b と cab の類似度は \(0\)

となります。これを順に並べた配列を LCP Array と言います。今回の例では \((0,1,1,0,0)\) となります。LCP Array は、Suffix Array が計算済みであれば適切なアルゴリズムにより \(O(|S|)\) で構築することができます。以下ではこの配列を \(LCPA\) と表します。

suffix の類似度の計算

3つの文字列 \(A,B,C\) が辞書順で \(A < B < C\) であるとき \(f(A,C)=\min(f(A,B),f(B,C))\) が成り立つことを利用すると、\(i<j\) のとき

\(f(S_{SA_i},S_{SA_j})=\min(f(S_{SA_i},S_{SA_{i+1}}),\ldots,f(S_{SA_{j-1}},S_{SA_j})) =\min(LCPA_i,\ldots,LCPA_{j-1})\)

となることから、2つの suffix の類似度は (2つの suffix が同一である場合を除いて) LCP Array の range min として得られることがわかります。

ここまでのまとめ

ここまでの結果を整理してみましょう。
\(SA_{x}=k\) であるとき、\(f(S_k,S_1)+\ldots+f(S_k,S_N)\) は項の順序を適当に並び替えることで、\(f(S_{SA_{x}},S_{SA_{1}})+\ldots+f(S_{SA_{x}},S_{SA_{N}})\) となります。これを次のように \(3\) つの部分に分けます。

• \(f(S_{SA_{x}},S_{SA_{1}})+\ldots+f(S_{SA_{x}},S_{SA_{x-1}})\)
• \(f(S_{SA_{x}},S_{SA_{x}})\)
• \(f(S_{SA_{x}},S_{SA_{x+1}})+\ldots+f(S_{SA_{x}},S_{SA_{N}})\)

\(2\) 番目は明らかに \(N-k\) なので残りの \(2\) つが求められればよいです。ここで、\(f(S_{SA_{i}},S_{SA_{j}})\) は LCP Array の range min として得られるということを前節で示しました。

したがって、元の問題を解くためには次の問題が解ければよいことがわかりました。

問題：
長さ \(N\) の数列 \(A=(A_0,\ldots,A_{N-1})\) が与えられる。\(A_0=0\) である。
\(i<j\) に対し \(g(i,j)=\min(A_i,A_{i+1},\ldots,A_{j-1})\) とするとき、\(k=1,\ldots,N\) のそれぞれについて、
\(g(1,k)+g(2,k)+\ldots+g(k-1,k)\) 及び \(g(k,k+1)+g(k,k+2)+\ldots+g(k,N)\)
を求めよ。

（この問題における \(k\) は、元の問題の \(SA_k\) にあたることに注意してください）

帰着された数列の問題を解く

\(g(1,k)+g(2,k)+\ldots+g(k-1,k)\) を \(B_k\) とし、以下では \(B_k\) の求め方を説明します。もう一方も同様に求めることができます。

\(k=1\) から順に \(B_k\) を求めることにします。\(B_k\) から、\(B_{k+1}\) を求めることを考えましょう。

\(\begin{array}{} B_k&=&g(1,k)&+&\ldots&+&g(k-1,k)&\\ B_{k+1}&=&g(1,k+1)&+&\ldots&+&g(k-1,k+1)&+&g(k,k+1)\\ &=& \min(g(1,k),A_k) &+&\ldots&+&\min(g(k-1,k),A_k)&+&A_k \end{array}\)

であることに注意すると、\(B_k\) は次のように求めることができます。

• 空の多重集合 \(X\) を用意する。\(k=1,\ldots,N\) の順に以下 \(3\) つの操作を行う
  • 1. \(X\) の全ての要素 \(x\) を \(\min(x,A_{k-1})\) に置き換える
  • 2. \(X\) に \(A_{k-1}\) を追加する
  • 3. \(X\) の全ての要素の和を求める。これが \(B_k\) である

これを愚直に実装すると最悪ケースで \(O(N^2)\) かかってしまいます。そこでかわりに、値と個数の組を管理することにします。操作 1 で実際に書き換えが発生するごとに \(X\) 内の数の種類数が減るため、操作を行う回数は全体で \(O(N)\) となり、priority queue などを用いることで最悪計算量 \(O(N \log N)\) でこの問題が解けました。

さらに、操作 2 で追加される数がその時点での \(\max X\) 以上であることを用いると、\(X\) を管理するデータ構造として単に stack を使っても全ての操作を行うことができ、この場合の計算量は \(O(N)\) となります。

以上により元の問題は \(O(N)\) で解けました。

なお、Suffix Array 及び LCP Array の構築は ACL (AtCoder Library) にあるため、C++を使用している場合はこれを用いると便利です。