AtCoderの文字列一致をどう選ぶ?KMP法とZ-algorithmを実装で理解

読了 約17分 たびすけ
AtCoderの文字列一致を学ぶ図。KMP法とZ-algorithmの考え方

次に読む記事

関連するテーマの記事を、先に確認できます。

問題文の形から、まず手法を選ぶ

AtCoderの文字列問題で迷ったら、最初に「何を何度比べる問題か」を書き出します。短い文字列を少数の位置で比べるだけなら単純比較で十分です。同じパターンを長い文字列の多くの位置で探すならKMP法、ひとつの文字列の各位置と先頭の共通部分を調べるならZ-algorithmが候補になります。

問題文の形 選ぶ手法 確認すること
短い範囲を少数回だけ比較する 単純比較 最悪時の比較回数が制約に収まるか
ひとつのパターンを長いテキストから探す KMP法 不一致時に、パターンの接頭辞を次の候補へ戻せるか
各位置と文字列の先頭を比べる Z-algorithm 共通接頭辞の長さを位置ごとに再利用できるか
一致した位置から一定幅で次の位置へ進む 一致判定+状態遷移 未一致で止まるか、訪問済みの状態へ戻るか

KMP法とZ-algorithmはどちらも比較のやり直しを減らしますが、出力する情報が違います。KMP法は「パターンをテキストへ走査したときの現在の一致長」を持ち、Z-algorithmは「各位置から先頭へ何文字一致するか」を配列にします。回転や無限列では、この一致結果の後にどの状態へ進むかまで追います。

添字と計算量を先にそろえる

以下のコードは文字列の先頭を0番目とする0-indexで書きます。長さNの文字列の添字は0からN - 1までです。長さMのパターンを開始位置rで比べる部分はtext[r:r + M]です。

各位置でパターンを最初から比べ直す素朴な方法は、最悪時にO(NM)になります。制約が大きいときは、既に一致した文字を次の比較へ渡す仕組みを用意します。KMP法とZ-algorithmは、前処理と走査を合わせてO(N + M)または対象長に対する線形時間で処理できます。

KMP法は、不一致時の戻り先を保存する

パターンの先頭からq文字が一致した状態で不一致になったとします。ここで、パターンの先頭と末尾が一致する最長の長さを次の候補にすれば、比較をやり直さずに済みます。その長さを接頭辞関数pi[q - 1]に保存します。パターン全体そのものは候補に含めません。

パターンababの接頭辞関数は次の通りです。最後まで一致した状態でも、先頭のabが末尾のabと重なっているので、次の候補は0ではなく2です。

位置 0 1 2 3
文字 a b a b
pi 0 0 1 2

abababを走査すると、添字3のb(bの2個目)で最初の一致が完成します。その直後にj = pi[j - 1] = 2へ戻るため、次のaとbで開始位置2の重なった一致も記録できます。

テキスト添字 文字 処理後のj 意味
0 a 1 先頭1文字が一致
1 b 2 先頭2文字が一致
2 a 3 先頭3文字が一致
3 b 4→2 一致を記録し、重なりへ戻る
4 a 3 次の一致候補を再利用
5 b 4→2 開始位置2の一致を記録

実装では、不一致の間だけj = pi[j - 1]を繰り返します。一致してパターン長に達した後も同じ戻り先へ移ることが、重複一致を落とさないポイントです。

def prefix_function(p):
    pi = [0] * len(p)
    for i in range(1, len(p)):
        j = pi[i - 1]
        while j and p[i] != p[j]:
            j = pi[j - 1]
        if p[i] == p[j]:
            j += 1
        pi[i] = j
    return pi

def kmp_starts(text, pattern):
    pi = prefix_function(pattern)
    starts = []
    j = 0
    for i, ch in enumerate(text):
        while j and ch != pattern[j]:
            j = pi[j - 1]
        if ch == pattern[j]:
            j += 1
        if j == len(pattern):
            starts.append(i - len(pattern) + 1)
            j = pi[j - 1]
    return starts

ABC135 Fで、一致位置を状態遷移へつなぐ

ABC135 F Strings of Eternityは、英小文字のsとtについて、tを何個連結した文字列が、sを無限に繰り返した列の部分文字列になるかを問います。|s|と|t|はそれぞれ5×105以下なので、全開始位置で比較をやり直す方法は使えません。

まずsをn + m - 1文字ぶん繰り返してKMP法を走査します。n = |s|、m = |t|とすると、開始位置0からn - 1までの各位置で長さmの一致を最後まで確認できます。得られた一致位置をmatched[r]に保存します。

一致位置rから次のtを置く開始位置は、周期をまたぐので (r + m) % n です。これは各頂点の出先が一つだけの状態遷移になります。

開始位置r tの一致 次の位置
0 一致 (0 + 2) % 7 = 2(未一致)
3 一致 (3 + 2) % 7 = 5(一致)
5 一致 (5 + 2) % 7 = 0(一致)

公式サンプル1のs = abcabab、t = abでは、一致する開始位置は0, 3, 5です。3→5→0→2と進むため、開始位置3からはtを3個つなげられ、出力は3になります。未一致へ出たら鎖は終わり、たどっている途中で訪問中の位置へ戻ったら一致が無限に続くので-1です。一致位置が一つもない場合は、空文字列の個数i = 0が答えになります。

import sys

def prefix_function(p):
    pi = [0] * len(p)
    for i in range(1, len(p)):
        j = pi[i - 1]
        while j and p[i] != p[j]:
            j = pi[j - 1]
        if p[i] == p[j]:
            j += 1
        pi[i] = j
    return pi

def solve():
    s = sys.stdin.readline().strip()
    t = sys.stdin.readline().strip()
    n, m = len(s), len(t)
    pi = prefix_function(t)
    matched = [False] * n
    j = 0
    for i in range(n + m - 1):
        ch = s[i % n]
        while j and ch != t[j]:
            j = pi[j - 1]
        if ch == t[j]:
            j += 1
        if j == m:
            start = i - m + 1
            if 0 <= start < n:
                matched[start] = True
            j = pi[j - 1]

    state = [0] * n
    dp = [0] * n
    answer = 0
    for start in range(n):
        if not matched[start] or state[start] == 2:
            continue
        path = []
        v = start
        base = 0
        while True:
            if not matched[v]:
                break
            if state[v] == 1:
                print(-1)
                return
            if state[v] == 2:
                base = dp[v]
                break
            state[v] = 1
            path.append(v)
            v = (v + m) % n
        while path:
            u = path.pop()
            base += 1
            dp[u] = base
            state[u] = 2
        answer = max(answer, dp[start])
    print(answer)

solve()

このコードではstate = 1が現在の探索中、state = 2が鎖の長さを確定済みという意味です。n + m - 1のKMP走査、一致位置の配列化、状態遷移の走査を合わせて計算量はO(n + m)、追加メモリもO(n + m)です。

公式サンプル 出力 確認する境界
abcabab / ab 3 重なった一致を鎖として数える
aa / aaaaaaa -1 状態遷移の閉路を検出する
aba / baaab 0 一致なしと-1を区別する

Z-algorithmは、各位置の一致長を再利用する

文字列Xの位置iから始まる接尾辞と、X全体の先頭が何文字一致するかをz[i]に記録します。z[0]は|X|です。対象の先頭へパターンを置いて後ろへ連結すれば、対象の各位置でパターンと一致する長さをZ値から読み取れます。AtCoder Libraryのz_algorithmの仕様も、各位置の最長共通接頭辞を線形時間で求める定義です。

右端まで一致を確認した区間を[left, right]として持ちます。新しい位置iがこの区間内なら、対応する既計算値をmin(right - i + 1, z[i - left])で初期値にできます。rightを越える部分だけをwhileで比較するため、同じ文字を何度も調べません。

def z_algorithm(s):
    n = len(s)
    z = [0] * n
    z[0] = n
    left = right = 0
    for i in range(1, n):
        if i <= right:
            z[i] = min(right - i + 1, z[i - left])
        while i + z[i] < n and s[z[i]] == s[i + z[i]]:
            z[i] += 1
        if i + z[i] - 1 > right:
            left, right = i, i + z[i] - 1
    return z

ABC430 Eで回転をZ値で判定する

ABC430 E Shift Stringは、同じ長さの二進文字列Aを左へ何回かシフトしてBに一致させる最小回数を求めます。公式解説では、B + A + AへZ-algorithmを適用する方法が示されています。

Aをk回左へシフトした文字列はA[k:] + A[:k]です。これはA + Aのk番目から始まる長さnの部分文字列なので、C = B + A + Aを作ります。Cの位置n + kから始まる接尾辞と先頭のBがn文字一致する条件は、z[n + k] >= nです。

公式サンプルの最初の組A = 1010001、B = 1000110で追います。n = 7としてk = 0は先頭2文字だけ一致し、k = 1は先頭文字から不一致、k = 2は7文字すべて一致します。したがって、z[7] = 2、z[8] = 0、z[9] >= 7となり、最小シフトは2です。

k 比較する回転 判定
0 1010001 と 1000110 z[7] = 2で不一致
1 0100011 と 1000110 z[8] = 0で不一致
2 1000110 と 1000110 z[9] >= 7で一致
import sys

def z_algorithm(s):
    n = len(s)
    z = [0] * n
    z[0] = n
    left = right = 0
    for i in range(1, n):
        if i <= right:
            z[i] = min(right - i + 1, z[i - left])
        while i + z[i] < n and s[z[i]] == s[i + z[i]]:
            z[i] += 1
        if i + z[i] - 1 > right:
            left, right = i, i + z[i] - 1
    return z

def solve_shift():
    t = int(sys.stdin.readline())
    for _ in range(t):
        a = sys.stdin.readline().strip()
        b = sys.stdin.readline().strip()
        n = len(a)
        z = z_algorithm(b + a + a)
        answer = -1
        for k in range(n):
            if z[n + k] >= n:
                answer = k
                break
        print(answer)

solve_shift()

kを0からn - 1まで調べるため、最初に一致した値が最小回数です。最初から同じならk = 0、どの回転でも一致しなければ-1です。公式サンプル全体の出力は順に2、-1、0、-1、22です。B + A + Aの長さは3nなので、各テストケースはO(n)時間・O(n)メモリで処理できます。

提出前に見る境界

場面 コードで見る場所 間違えやすい点
KMPの不一致 j = pi[j - 1] j = 0まで戻る前に候補を使い切る
KMPの重複一致 一致直後のj = pi[j - 1] 一致後に0へ戻して次の重なりを失わない
ABC135 Fの未一致 matchedと鎖の終端 一致なしの答え0と、閉路の-1を混同しない
ABC135 Fの周期 (r + m) % n 文字列の位置ではなく周期上の状態を進める
ABC430 Eの回転 z[n + k] >= n k = 0を含め、n回目を重ねて数えない

問題文を読んだら、単純比較で制約に収まるかを確認し、収まらないときにKMP法かZ-algorithmへ切り替えます。KMP法では戻り先と重複一致、Z-algorithmではZ-boxと添字n + kを小さな入力で紙に書くと、式とコードの対応を保ったまま実装できます。そこからABC135 Fの状態遷移、ABC430 Eの回転判定へ進めば、同じ線形走査の考え方を別の問題へ移せます。