Day 074-Q1 — Aho-Corasick + DP on Automaton（禁止パターン回避・行列累乗）

問題

$M$ 個の「禁止パターン」$P_1, P_2, \ldots, P_M$（英小文字）が与えられる。長さちょうど $N$ の英小文字列のうち、どの禁止パターンも部分文字列として含まない文字列の個数を $998244353$ で割った余りを求めよ。

入出力例

5 2
ab
bc

6696576

3 1
aaa

17549

概念図: Aho-Corasick オートマトンと状態遷移行列

ヒント

# Aho-Corasick の fail 伝播（BFS）
q = deque()
for c in range(26):
    if goto_list[0][c] == -1:
        goto_list[0][c] = 0  # ルートの未登録 → ルート自身へ
    else:
        fail[goto_list[0][c]] = 0
        q.append(goto_list[0][c])

while q:
    u = q.popleft()
    if output[fail[u]]:
        output[u] = True  # fail 先が BAD → 自分も BAD
    for c in range(26):
        v = goto_list[u][c]
        if v == -1:
            goto_list[u][c] = goto_list[fail[u]][c]
        else:
            fail[v] = goto_list[fail[u]][c]
            q.append(v)

模範解答

import sys
from collections import deque
input = sys.stdin.readline

MOD = 998244353

def solve():
    N, M = map(int, input().split())
    patterns = [input().strip() for _ in range(M)]

    ALPHA = 26
    goto_list = [[-1] * ALPHA]
    fail = [0]
    output = [False]

    # Trie 挿入
    for pat in patterns:
        cur = 0
        for ch in pat:
            c = ord(ch) - ord('a')
            if goto_list[cur][c] == -1:
                goto_list[cur][c] = len(goto_list)
                goto_list.append([-1] * ALPHA)
                fail.append(0)
                output.append(False)
            cur = goto_list[cur][c]
        output[cur] = True

    # BFS で fail リンク & goto 完全化
    q = deque()
    for c in range(ALPHA):
        if goto_list[0][c] == -1:
            goto_list[0][c] = 0
        else:
            fail[goto_list[0][c]] = 0
            q.append(goto_list[0][c])

    while q:
        u = q.popleft()
        if output[fail[u]]:
            output[u] = True
        for c in range(ALPHA):
            v = goto_list[u][c]
            if v == -1:
                goto_list[u][c] = goto_list[fail[u]][c]
            else:
                fail[v] = goto_list[fail[u]][c]
                q.append(v)

    # 安全状態で遷移行列構築
    safe = [i for i in range(len(goto_list)) if not output[i]]
    idx = {s: i for i, s in enumerate(safe)}
    K = len(safe)

    T = [[0] * K for _ in range(K)]
    for s in safe:
        for c in range(ALPHA):
            nxt = goto_list[s][c]
            if not output[nxt]:
                T[idx[s]][idx[nxt]] = (T[idx[s]][idx[nxt]] + 1) % MOD

    # 行列累乗
    def mat_mul(A, B):
        n = len(A)
        C = [[0] * n for _ in range(n)]
        for i in range(n):
            for k in range(n):
                if A[i][k] == 0:
                    continue
                for j in range(n):
                    C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % MOD
        return C

    def mat_pow(Mat, p):
        n = len(Mat)
        R = [[int(i == j) for j in range(n)] for i in range(n)]
        while p:
            if p & 1:
                R = mat_mul(R, Mat)
            Mat = mat_mul(Mat, Mat)
            p >>= 1
        return R

    res = mat_pow(T, N)
    init = idx[0]
    print(sum(res[init]) % MOD)

solve()

Step-by-Step 解説

Step 1: Aho-Corasick の構築

複数パターンを同時に検索するデータ構造。Trie + failure link（KMP の失敗関数の一般化）で構成される。BFS で各状態の遷移を「完全関数」化する（どの文字を入力しても次の状態が一意に定まる）。

output[u] = True は「状態 u に到達 = 禁止パターンが現れた」を意味し、fail リンクを伝播させる。

Step 2: 安全状態と遷移行列

output[u] = False の状態のみを残し、$K \times K$ の遷移行列 $T$ を作る。$T[i][j]$ = 状態 $i$ から 1 文字追加したとき状態 $j$ に遷移する文字の種類数。

Step 3: 行列累乗

長さ $N$ の文字列 → $N$ 回の遷移 → $T^N$ の (0→j) 成分の総和が答え。$N \le 10^{18}$ なので繰り返し二乗法で $O(K^3 \log N)$。

Step 4: 計算量

• Aho-Corasick 構築: $O(\sum |P| \times 26)$
• 行列累乗: $O(K^3 \log N)$ ここで $K \le 201$
• 全体: 約 $201^3 \times 60 \approx 4.8 \times 10^8$（PyPy 推奨）

よくあるミス

次のステップ

• 発展問題: 禁止パターンを「正規表現」に拡張 → NFA → DFA 変換 → 行列累乗
• 類題: AtCoder ABC DP 問題「禁止文字列を含まない数え上げ」
• さらに: Suffix Automaton を使った $O(|P|)$ クエリへの昇華

自己評価