Day 103-Q2 | 桁DP × Aho-Corasickオートマトン

問題

整数 $N$（10進表記）と $K$ 個の禁止パターン文字列 $P_1,\dots,P_K$（数字のみ）が与えられる。$1$ 以上 $N$ 以下の整数のうち、その10進表記が $P_1,\dots,P_K$ のいずれも部分文字列として含まないものの個数を求めよ。

複数パターンを同時に回避するには Aho-Corasickオートマトン が必要。$K$個のパターンから構築したTrieに fail リンクを張ってオートマトン化し、各状態が「いずれかのパターンで終わっている（かfailを辿った先がそう）」かを表すdangerフラグを事前計算する。桁DPの状態に「現在のオートマトンのノード番号」を追加し、遷移先がdangerな遷移をすべて禁止することで数え上げる。

入力形式

N K
P_1
P_2
:
P_K

制約

$1 \le N \le 10^{18}$

$1 \le K \le 10$

各$P_i$は長さ$1$〜$6$の数字列

数字のみで構成

入出力例

入力例1

100 2
13
4

出力例1

$1$〜$100$のうち"13"または"4"を含む数は20個あり、$100-20=80$。

入力例2

1000000000000000000 3
13
4
666

出力例2

118730162908827870

$N=10^{18}$の巨大ケースでも、状態数は$\sum|P_i|+1$程度で桁DPが高速に完了する。

概念図

ヒント（段階的開示）

ヒント1: 方向性

$1$から$N$まで愚直にループして部分文字列判定するのは$O(N\cdot\sum|P_i|)$で$N\le10^{18}$では全く間に合わない。「上から$i$桁目まで確定した」状態を持つ桁DPが必要だが、単純な状態だけでは「禁止パターンを含んでいないか」を判定できない。

ヒント2: アプローチ

複数パターンの「今どこまでマッチしているか」を管理する構造がAho-Corasickオートマトン。$K$個のパターンからTrieを作りfailリンク（KMPの一般化）を張ることで、任意の状態から任意の文字への遷移が$O(1)$で引ける完全なgoto関数を作る。各状態について「そこ、またはfailを辿った先のどれかがパターン終端か」を表すdangerフラグを伝播させ、桁DPの遷移でdanger状態への遷移を禁止するだけでよい。

ヒント3: 誘導（コード骨格）

def rec(pos, node, tight, started):
    if danger[node]:
        return 0
    if pos == length:
        return 1
    limit = int(n_str[pos]) if tight else 9
    total = 0
    for d in range(limit + 1):
        n_started = started or d != 0
        n_node = node
        if n_started:
            n_node = goto[node][str(d)]
            if danger[n_node]:
                continue
        total += rec(pos+1, n_node, tight and d==limit, n_started)
    return total

started（先頭の0埋めを終えたか）を管理し、先頭の0がオートマトンの状態遷移に影響しないようにする点がポイント。

模範解答 (Python)

import sys
from collections import deque


def build_automaton(patterns):
    trie_goto = [dict()]
    is_end = [False]
    for p in patterns:
        cur = 0
        for ch in p:
            if ch not in trie_goto[cur]:
                trie_goto.append(dict())
                is_end.append(False)
                trie_goto[cur][ch] = len(trie_goto) - 1
            cur = trie_goto[cur][ch]
        is_end[cur] = True

    size = len(trie_goto)
    fail = [0] * size
    danger = is_end[:]
    goto = [dict(trie_goto[i]) for i in range(size)]

    for ch in '0123456789':
        if ch not in goto[0]:
            goto[0][ch] = 0

    q = deque()
    for ch, nxt in trie_goto[0].items():
        fail[nxt] = 0
        q.append(nxt)

    while q:
        u = q.popleft()
        danger[u] = danger[u] or danger[fail[u]]
        for ch in '0123456789':
            if ch in trie_goto[u]:
                v = trie_goto[u][ch]
                fail[v] = goto[fail[u]][ch]
                goto[u][ch] = v
                q.append(v)
            else:
                goto[u][ch] = goto[fail[u]][ch]

    return goto, danger


def count_valid(n_str, patterns):
    goto, danger = build_automaton(patterns)
    length = len(n_str)
    memo = {}

    def rec(pos, node, tight, started):
        if danger[node]:
            return 0
        if pos == length:
            return 1
        key = (pos, node, tight, started)
        if key in memo:
            return memo[key]

        limit = int(n_str[pos]) if tight else 9
        total = 0
        for d in range(limit + 1):
            n_started = started or d != 0
            n_node = node
            if n_started:
                n_node = goto[node][str(d)]
                if danger[n_node]:
                    continue
            total += rec(pos + 1, n_node, tight and d == limit, n_started)
        memo[key] = total
        return total

    return rec(0, 0, True, False) - 1  # 0を除く


def solve():
    data = sys.stdin.read().split()
    n_str = data[0]
    k = int(data[1])
    patterns = data[2:2 + k]
    print(count_valid(n_str, patterns))


solve()

計算量: $O(|N| \times \text{状態数} \times 4)$（桁数18・Aho-Corasick状態数$\le\sum|P_i|+1$・tight/started各2通り）。

Step-by-Step 解説

1Trie構築
$K$個のパターンを1本のTrieにまとめ、各ノードに`is_end`を記録する。

2failリンクとgoto関数（BFS）
KMPを一般化したfailリンクをBFSで構築し、任意の文字への完全な遷移表を作る。

3dangerフラグの伝播
`danger[u] = danger[u] or danger[fail[u]]`をBFS順に伝播させ、間接的にパターンを含む状態も検出する。

4桁DPでの数え上げ
(桁位置, オートマトン状態, tight, started)を状態とし、danger遷移を除外して数える。

よくあるミス

ミス	原因	正しい書き方
パターンごとに別々のオートマトンを使う	合成が複雑になり計算量も増える	1本のTrie+failリンクで統合したAho-Corasickオートマトンを構築する
ルートの`goto`に未設定文字が残りKeyError	ルートはfail先の特別扱いが必要	ルートの存在しない文字の遷移はあらかじめ自分自身に設定しておく
`started`フラグを持たず先頭0をそのまま流し込む	"004"のような表記されない先頭0が判定に混ざる	数字が始まるまで`node`を更新せず、開始した瞬間から遷移させる
答えに0を含めてしまう	`rec`は0以上$N$以下の個数を返す	最終的に`-1`して0を除外する

次のステップ

発展: パターンに`?`（任意の1文字）のようなワイルドカードを許可する拡張
発展: 「ちょうど$M$回出現する数」を数えるバージョン（DP状態に出現回数の次元を追加）
次回予告: 木上のMo's Algorithm（オイラーツアー圧縮 + パス上の色の種類数クエリ）

自己評価

完全理解おおむね理解要復習

Day 103-Q2 — 桁DP × Aho-Corasickオートマトン

問題