Aho–Corasick
パターンを Trie に入れ、各ノードに「そのノードの文字列の真の接尾辞のうち、Trie にある最長のもの」への失敗リンクを張ります。テキストを 1 文字ずつ読みながら Trie をたどり、進めないときは失敗リンクで戻ることで、すべてのパターンの出現をまとめて数えます。文字は辞書のキーにできれば何でもよいです(整数のリストも使えます)。
以下、 $\lvert T \rvert$ はテキストの長さ、 $\sum \lvert P \rvert$ はパターンの長さの合計です。
AhoCorasick(patterns): パターン(空でない文字列)のリストから構築。 $O(\sum \lvert P \rvert)$count(text): 各パターンがテキストに現れる回数(重なりも数える)のリスト。 $O(\lvert T \rvert + \sum \lvert P \rvert)$step(v, c): 状態 $v$ で文字 $c$ を読んだあとの状態。何も読んでいない状態は $0$。テキストを先頭から読むときはならし $O(1)$out[v]: 状態 $v$ までに読んだ文字列の接尾辞になっているパターンの個数node[i]: パターン $i$ を読み終えた状態
step で求めた遷移は覚えておくので、文字列についての DP で同じ遷移を何度も使っても速いです。
例(パターンを 1 つも含まない長さ $L$ の英小文字列の個数を $998244353$ で割った余り):
ac = AhoCorasick(patterns)
dp = {0: 1}
for _ in range(L):
ndp = {}
for v, x in dp.items():
for c in "abcdefghijklmnopqrstuvwxyz":
to = ac.step(v, c)
if ac.out[to] == 0:
ndp[to] = (ndp.get(to, 0) + x) % 998244353
dp = ndp
ans = sum(dp.values()) % 998244353
コード
class AhoCorasick:
def __init__(self, patterns):
self.go = [{}] # go[v][c] = 状態 v で文字 c を読んだあとの状態(Trie の辺と、step で求めたもの)
self.node = [] # node[i] = パターン i を読み終えた状態
for s in patterns:
v = 0
for c in s:
to = self.go[v].get(c)
if to is None:
to = len(self.go)
self.go[v][c] = to
self.go.append({})
v = to
self.node.append(v)
n = len(self.go)
self.fail = [0] * n # fail[v] = v の文字列の真の接尾辞で、Trie にある最長のもの
self.out = [0] * n # out[v] = v の文字列の接尾辞になっているパターンの個数
for v in self.node:
self.out[v] += 1
self.order = [0] # BFS 順
for v in self.order:
for c, to in self.go[v].items():
f = self.fail[v]
while f and c not in self.go[f]:
f = self.fail[f]
if v and c in self.go[f]:
self.fail[to] = self.go[f][c]
self.out[to] += self.out[self.fail[to]]
self.order.append(to)
def step(self, v, c):
# 状態 v で文字 c を読んだあとの状態
go = self.go
path = []
while c not in go[v]:
path.append(v)
if v == 0:
go[0][c] = 0
break
v = self.fail[v]
res = go[v][c]
for u in path:
go[u][c] = res
return res
def count(self, text):
# 各パターンが text に現れる回数のリスト
cnt = [0] * len(self.go)
v = 0
for c in text:
v = self.step(v, c)
cnt[v] += 1
for v in reversed(self.order[1:]):
cnt[self.fail[v]] += cnt[v]
return [cnt[v] for v in self.node]