AtCoder ABC353 EをTrieで解く：最長共通接頭辞の合計

ABC353 E – Yet Another Sigma Problemは、文字列を2つずつ比べたときの最長共通接頭辞（先頭から一致し続ける部分、LCP）の長さを、すべての組について合計する問題です。Trieに文字列を順番に挿入し、各prefixのnodeを以前に通った文字列数を加えると、この合計を数えられます。足すのは現在の文字列を通す前の数です。

全ペアを直接調べると約450億組になる

公式問題の制約は、文字列数Nが最大3×10⁵、文字数の合計Tも最大3×10⁵です。Nが最大のとき、比較する組はN(N − 1) ÷ 2、約450億組になります。1組ずつLCPを求める方法は、文字を比べ始める前に組の数だけで重くなります。

AtCoder公式解説はTrieを使い、同じprefixを持つ文字列が共有するnodeに、そのprefixを通過した文字列数を記録します。Trieでは文字列を先頭からたどるため、深さdのnodeは長さdのprefixを表します。同じprefixを持つ文字列は、そのnodeを共通して通ります。

nodeを通った既出文字列数を足す

passed[v]を、これまでに挿入した文字列のうちnode vを通った本数とします。新しい文字列の各文字をTrieでたどるとき、その文字の先のnodeに着いたら、まずpassed[v]を答えへ加え、その後でpassed[v]を1増やします。これで加算時の数は、現在の文字列より先に挿入された文字列だけを表します。

公式サンプルのab、abc、arcを追う

公式サンプル1の入力例

3
ab abc arc

出力例1

公式サンプル1では、abを挿入した時点で、node aとnode abのpassedがそれぞれ1になります。次のabcは、node aで1、node abで1を足し、初めて通るnode abcでは0を足します。この文字列が加える値は合計2です。

最後のarcは、最初の文字aで、そこを通った既出文字列2本分を足します。続くrとcには既出の文字列が通っていないので、この文字列の加算は2です。答えは2 + 2 = 4となり、公式出力と一致します。各ペアを直接見ると、LCPはabとabcで2、abとarcで1、abcとarcで1です。

標準入力から答えを出すPythonコード

子nodeは文字からnode番号への辞書で持ちます。必要なnodeがなければ作り、現在の文字列が通るたびに既出数を加えてから更新します。

import sys

data = sys.stdin.buffer.read().split()
n = int(data[0])
strings = data[1:1 + n]

children = [{}]
passed = [0]
answer = 0

for string in strings:
    node = 0
    for char in string:
        next_node = children[node].get(char)
        if next_node is None:
            next_node = len(children)
            children[node][char] = next_node
            children.append({})
            passed.append(0)

        node = next_node
        answer += passed[node]
        passed[node] += 1

print(answer)

各ペアのLCPをちょうど一度ずつ数える

すでに挿入した文字列tと、いま挿入する文字列sのLCPが長さLだとします。両方が通るnodeは、深さ1からLまでのL個です。各nodeのpassedにはtが1本含まれるため、sを挿入するときにtとの共通prefixが1文字分ずつ、合計Lだけ加わります。最初に異なる文字へ進んだ後は共有nodeがなく、加算されません。

新しい文字列は、それより前の文字列すべてとの組を一度ずつ数えます。どの組も後から挿入される側の文字列でだけ数えるので、全組のLCP長の合計になります。各通過nodeのpassedを加算後に1増やす操作は、その文字列も以後の文字列に対する既出文字列へ含めるため、同じ数え方を保ちます。

重複文字列とprefix関係も同じ規則で扱える

同じ文字列が複数回現れる場合、後から挿入した文字列は同じnode列をたどり、各nodeで以前の同じ文字列の本数を加えます。したがって、重複した2文字列のLCPも文字列長ぶん数えられます。

一方がもう一方のprefixなら、短い文字列はその長さのnodeまでしか進みません。そこで共有nodeの数も短い方の長さで止まります。終端を示すnodeや別の終端カウンタは不要です。rootは空のprefixで長さ0のため加算せず、1文字の文字列もその1文字のnodeだけを数えます。

計算量は期待O(T)、記憶量はO(T)

各文字につき子nodeの検索を行い、nodeがなければ1つ作ります。Pythonの辞書操作が期待O(1)時間であることを前提にすると、総文字数Tに対して期待時間計算量はO(T)です。作られるnodeと辺はそれぞれ高々T個で、passedを含むTrieの記憶量はO(T)です。入力文字列の保持を含めてもO(T)の範囲です。

複数文字列のprefixをまとめて数える課題から、1本の文字列の中でパターンの一致位置を探す課題へ進む場合は、KMP法とZ-algorithmの実装解説が次の題材になります。