字符串算法

string_search.py

def prefix_function(p: str) -> list[int]:
    """pi[i] = length of the longest proper border of p[:i + 1]."""
    pi = [0] * len(p)
    k = 0
    for i in range(1, len(p)):
        while k > 0 and p[i] != p[k]:
            k = pi[k - 1]
        if p[i] == p[k]:
            k += 1
        pi[i] = k
    return pi


def kmp_search(text: str, pattern: str) -> list[int]:
    """Start positions of every (possibly overlapping) occurrence."""
    if not pattern:
        return list(range(len(text) + 1))
    pi = prefix_function(pattern)
    hits, k = [], 0
    for i, ch in enumerate(text):
        while k > 0 and ch != pattern[k]:
            k = pi[k - 1]
        if ch == pattern[k]:
            k += 1
        if k == len(pattern):
            hits.append(i - k + 1)
            k = pi[k - 1]
    return hits


if __name__ == "__main__":
    print(prefix_function("ababaca"))         # [0, 0, 1, 2, 3, 0, 1]
    print(kmp_search("abababcabab", "abab"))  # [0, 2, 7]

python string_search.py