Day 095-Q3 | Suffix Automaton拡張（辞書順K番目相異なる部分文字列）

問題

英小文字からなる文字列 $S$ が与えられる。$S$ の相異なる（重複を除いた）非空部分文字列を辞書順に並べたとき、$K$ 番目に小さいものを出力せよ。存在しない場合は -1 を出力せよ。

入力形式

S
K

制約

$1 \le |S| \le 200000$

$S$ は英小文字のみ

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

入出力例

入力例1

aba
3

出力例1

aba

相異なる部分文字列を辞書順に並べると a, ab, aba, b, ba の5個。3番目は aba。K=6のときは -1（存在しない）。

概念図

ヒント（段階的開示）

ヒント1: 方向性

「相異なる部分文字列」を全列挙すると $O(N^2)$ 個になり得るため、$N \le 2\times10^5$ では列挙不可能。部分文字列の集合をコンパクトに表現するデータ構造上で、辞書順に「何番目か」を数えながら1文字ずつ進む必要がある。

ヒント2: アプローチ

Suffix Automaton（SAM）を構築すると、$S$ の全ての部分文字列は「初期状態からある状態までの経路」と1対1に対応する。各状態 $v$ について「$v$ から先に進める相異なる経路の数」$dp[v]$（$v$で止まる1通り + 各遷移先のdpの総和）を、状態を len の降順に処理して計算する。初期状態から遷移先の文字を昇順に見ながらdpを使って $K$ 番目の経路をそのまま辿ればよい。

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

order = sorted(range(len(sa)), key=lambda i: sa[i].len, reverse=True)
dp = [1] * len(sa)
for v in order:
    for c, to in sa[v].next.items():
        dp[v] += dp[to]

total = dp[0] - 1   # 初期状態の"止まる"は空文字列なので除く
if K > total:
    print(-1)
else:
    v, res = 0, []
    while True:
        if v != 0:
            if K == 1:
                break
            K -= 1
        for c in sorted(sa[v].next):
            to = sa[v].next[c]
            if K <= dp[to]:
                res.append(c); v = to; break
            K -= dp[to]
    print(''.join(res))

模範解答 (Python)

import sys


class State:
    __slots__ = ['len', 'link', 'next']

    def __init__(self):
        self.len = 0
        self.link = -1
        self.next = {}


def build_sam(s):
    sa = [State()]
    last = 0
    for ch in s:
        cur = len(sa)
        sa.append(State())
        sa[cur].len = sa[last].len + 1
        p = last
        while p != -1 and ch not in sa[p].next:
            sa[p].next[ch] = cur
            p = sa[p].link
        if p == -1:
            sa[cur].link = 0
        else:
            q = sa[p].next[ch]
            if sa[p].len + 1 == sa[q].len:
                sa[cur].link = q
            else:
                clone = len(sa)
                sa.append(State())
                sa[clone].len = sa[p].len + 1
                sa[clone].next = dict(sa[q].next)
                sa[clone].link = sa[q].link
                while p != -1 and sa[p].next.get(ch) == q:
                    sa[p].next[ch] = clone
                    p = sa[p].link
                sa[q].link = clone
                sa[cur].link = clone
        last = cur
    return sa


def main():
    data = sys.stdin.read().split()
    s = data[0]
    k = int(data[1])

    sa = build_sam(s)
    n = len(sa)

    order = sorted(range(n), key=lambda i: sa[i].len, reverse=True)
    dp = [1] * n
    for v in order:
        for c, to in sa[v].next.items():
            dp[v] += dp[to]

    total = dp[0] - 1
    if k > total:
        print(-1)
        return

    v = 0
    res = []
    while True:
        if v != 0:
            if k == 1:
                break
            k -= 1
        for c in sorted(sa[v].next.keys()):
            to = sa[v].next[c]
            if k <= dp[to]:
                res.append(c)
                v = to
                break
            k -= dp[to]

    print(''.join(res))


main()

計算量: SAM構築 $O(N)$。dp計算 $O(N)$。K番目クエリの復元は解の長さ $L$ に対し $O(L\log\sigma)$（$\sigma\le26$）。全体で $O(N\log\sigma)$。

Step-by-Step 解説

1SAMの構築
$S$の全部分文字列を「初期状態からの経路」として圧縮表現する。状態数・遷移数は $O(N)$。

2経路数dpの計算
状態をlenの降順に処理し、$dp[v]=1+\sum dp[to]$を計算する。

3初期状態の特別扱い
初期状態で止まる選択肢は空文字列なので数えない。全体数は$dp[0]-1$。

4K番目の経路を辿る
遷移先の文字を昇順に見て$dp[to]$と$K$を比較しながら貪欲に降りていく。

よくあるミス

ミス	原因	正しい書き方
初期状態でも「止まる」選択肢を数えてしまう	空文字列を非空部分文字列として誤カウント	`if v != 0`で停止判定をスキップする
dpを`len`の昇順で計算してしまう	依存関係の向きを取り違え	遷移は必ずlenが増える方向なので降順に計算
遷移の子を文字順にソートせず辞書の挿入順で見てしまう	dictのキー順は文字コード順とは限らない	`sorted(sa[v].next.keys())`で明示的にソート
Kの範囲チェックを怠り無限ループになる	総数を超えるKの事前チェック漏れ	walk前に`k > dp[0]-1`なら-1を返す

次のステップ

発展: 複数文字列に対する一般化SAM上でのK番目共通部分文字列クエリ
発展: オンラインクエリ（複数のKが与えられる場合）は同じdpを使い回せる
次回予告: Matrix-Forest定理（全域森数え上げ）

自己評価

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

Day 095-Q3 — Suffix Automaton拡張（辞書順K番目相異なる部分文字列）

問題