Official

G - Count Substring Query Editorial by toam


~ をどの英小文字よりも辞書順で大きい文字とします.

各クエリで求めるものは \(S[1:],S[2:],\ldots,S[N:]\) のうち,prefix が \(T\) と一致するものの個数です.これは,\(T\) の末尾に ~ を付加した文字列を \(T'\) としたときに \(T\leq S[i:]\leq T'\) を満たす \(i\) の個数に等しいです.あらかじめ \(S[1:],S[2:],\ldots,S[N:]\) を辞書順に並べた列(Suffix Array)を求めておき,この列上で二分探索をすることで, \(T\leq S[i:]\leq T'\) を満たす \(i\) の個数を高速に求めることが出来ます.計算量は Suffix Array の構築が \(O(|S|)\),各クエリが \(O(|T|\log |S|)\) です.

from atcoder.string import suffix_array

S = input()
SA = suffix_array(S)

def binary_search(T):
    ng, ok = -1, len(S)
    while ok - ng > 1:
        mid = (ng + ok) // 2
        l = SA[mid]
        if T <= S[l : l + len(T)]:
            ok = mid
        else:
            ng = mid
    return ok

Q = int(input())
for _ in range(Q):
    T = input()
    print(binary_search(T + "~") - binary_search(T))

(おまけ)言語の標準機能を上手に利用すると簡潔に書けます.

from atcoder.string import suffix_array
import bisect

S = input()
SA = suffix_array(S)

Q = int(input())
for _ in range(Q):
    T = input()
    l = bisect.bisect_left(SA, T, key=lambda x: S[x : x + len(T)])
    r = bisect.bisect(SA, T, key=lambda x: S[x : x + len(T)])
    print(r - l)

posted:
last update: