Day 073-Q1 — Suffix Tree（Ukkonen's Algorithm・LCS・パターンマッチ）

問題

長さ $N$ の文字列 $S$（英小文字）と $Q$ 個のクエリが与えられる。各クエリは文字列 $P_i$（長さ $|P_i|$）で、$S$ 中に $P_i$ が出現する回数と、最初の出現位置（0-indexed, 左端）を最小のものを答えよ。

すべてのクエリ長の総和 $\sum |P_i| \le 2 \times 10^5$ であることが保証されている。

入出力例

10 3
abcabcabc
ab
abc
abcd

3 0
3 0
0 -1

概念図: Suffix Array と二分探索

ヒント

def lower(sa, s, p):
    m = len(p)
    lo, hi = 0, len(sa)
    while lo < hi:
        mid = (lo + hi) // 2
        if s[sa[mid]:sa[mid]+m] < p:
            lo = mid + 1
        else:
            hi = mid
    return lo

def upper(sa, s, p):
    m = len(p)
    lo, hi = 0, len(sa)
    while lo < hi:
        mid = (lo + hi) // 2
        if s[sa[mid]:sa[mid]+m] <= p:
            lo = mid + 1
        else:
            hi = mid
    return lo

模範解答

import sys
input = sys.stdin.readline

def build_sa(s):
    """Suffix Array O(N log^2 N) ダブリング法"""
    n = len(s)
    if n == 0:
        return []
    sa = list(range(n))
    rank = [ord(c) for c in s]
    k = 1
    while k < n:
        r = rank
        def key(i):
            return (r[i], r[i + k] if i + k < n else -1)
        sa.sort(key=key)
        new_rank = [0] * n
        new_rank[sa[0]] = 0
        for j in range(1, n):
            prev_key = (r[sa[j-1]], r[sa[j-1]+k] if sa[j-1]+k < n else -1)
            curr_key = (r[sa[j]],   r[sa[j]+k]   if sa[j]+k   < n else -1)
            new_rank[sa[j]] = new_rank[sa[j-1]] + (1 if curr_key != prev_key else 0)
        rank = new_rank
        if rank[sa[-1]] == n - 1:
            break
        k <<= 1
    return sa

def solve():
    data = sys.stdin.read().split()
    idx = 0
    N, Q = int(data[idx]), int(data[idx+1]); idx += 2
    S = data[idx]; idx += 1

    sa = build_sa(S)

    def lower(p):
        m = len(p)
        lo, hi = 0, len(sa)
        while lo < hi:
            mid = (lo + hi) // 2
            if S[sa[mid]:sa[mid]+m] < p:
                lo = mid + 1
            else:
                hi = mid
        return lo

    def upper(p):
        m = len(p)
        lo, hi = 0, len(sa)
        while lo < hi:
            mid = (lo + hi) // 2
            if S[sa[mid]:sa[mid]+m] <= p:
                lo = mid + 1
            else:
                hi = mid
        return lo

    out = []
    for _ in range(Q):
        P = data[idx]; idx += 1
        l = lower(P)
        u = upper(P)
        count = u - l
        if count == 0:
            out.append("0 -1")
        else:
            first = min(sa[l:u])
            out.append(f"{count} {first}")

    print('\n'.join(out))

solve()

Step-by-Step 解説

Step 1: Suffix Array の構築

Suffix Array は文字列のすべての接尾辞を辞書順にソートした際のインデックス列。ダブリング法では長さ $k$ のランクを用いて長さ $2k$ のランクを計算し、$O(N \log N)$ のソートを $O(\log N)$ 回繰り返す。全体計算量は $O(N \log^2 N)$。

Step 2: パターン検索の二分探索

Suffix Array 上で S[sa[i]:sa[i]+|P|] と P を比較しながら二分探索。

• lower(P): $S[\text{sa}[i]:]$ が $P$ 以上になる最初の位置
• upper(P): $S[\text{sa}[i]:]$ が $P$ より大きくなる最初の位置
• 出現回数 = upper - lower

各クエリの計算量は $O(|P| \log N)$。

Step 3: 最初の出現位置

sa[l:u] の最小値を取る。計算量は出現区間の幅に比例するが、問題の制約内では十分高速。

よくあるミス

次のステップ

• 発展: SA-IS $O(N)$ 構築を実装する
• 発展: LCP Array（Kasai's algorithm）+ Sparse Table で LCE クエリを $O(1)$ に
• 発展: Suffix Automaton (SAM) で出現回数を $O(|P|)$ に短縮

自己評価