Knuth-Morris-Pratt algorithm in Python
Linear-time matching using the prefix function; the text pointer never moves backward.
- File
knuth_morris_pratt.py, 67 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/knuth-morris-pratt/. The authors are listed in prov/F0V6W1.jsonl.
from collections.abc import Sequence
from typing import TypeVar
T = TypeVar("T")
def compute_prefix_function(p: Sequence[T]) -> list[int]:
"""Return pi, where pi[q] is the length of the longest proper border of p[:q+1].
Time O(m), at most 2(m - 1) equality tests.
"""
m = len(p)
pi = [0] * m
k = 0
for q in range(1, m):
while k > 0 and p[k] != p[q]:
k = pi[k - 1]
if p[k] == p[q]:
k += 1
pi[q] = k
return pi
def compute_next(p: Sequence[T]) -> list[int]:
"""Return the strong failure table nxt[0..m] of KMP, for a non-empty p.
For 0 <= q < m, nxt[q] is the length k of the longest proper border of
p[:q] with p[k] != p[q], or -1 when there is none. nxt[m] is pi[m - 1].
Time O(m).
"""
m = len(p)
pi = compute_prefix_function(p) # @line L12
nxt = [0] * (m + 1)
nxt[0] = -1 # @line L13
for q in range(1, m): # @line L14
k = pi[q - 1] # @line L15
nxt[q] = nxt[k] if p[k] == p[q] else k # @line L16
nxt[m] = pi[m - 1] # @line L17
return nxt # @line L18
def knuth_morris_pratt(text: Sequence[T], pattern: Sequence[T]) -> list[int]:
"""Return every 0-based shift s with text[s:s+m] == pattern, in increasing order.
General alphabet (only equality tests). Occurrences may overlap. The empty
pattern occurs at each shift 0..n; a pattern longer than the text has none.
The text is read once, left to right, never backward.
Time O(n + m), at most 2n comparisons in the search; extra space O(m).
"""
n, m = len(text), len(pattern) # @line L1
if m == 0: # @line L2
return list(range(n + 1))
nxt = compute_next(pattern) # @line L3
out: list[int] = []
q = 0 # @line L4
for i in range(n): # @line L5
c = text[i]
while q >= 0 and pattern[q] != c: # @line L6
q = nxt[q] # @line L7
q += 1 # @line L8
if q == m: # @line L9
out.append(i - m + 1) # @line L10
q = nxt[m] # @line L11
return out
test_knuth_morris_pratt.py 68 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/knuth-morris-pratt/. The authors are listed in prov/F0V6W1.jsonl.
import random
from knuth_morris_pratt import compute_next, compute_prefix_function, knuth_morris_pratt
def oracle(t: str, p: str) -> list[int]:
return [s for s in range(len(t) - len(p) + 1) if t.startswith(p, s)]
def border_oracle(p: str) -> list[int]:
return [max(k for k in range(q + 1) if p[:k] == p[q + 1 - k : q + 1]) for q in range(len(p))]
def next_oracle(p: str) -> list[int]:
m = len(p)
out = []
for q in range(m):
ks = [k for k in range(q) if p[:k] == p[q - k : q] and p[k] != p[q]]
out.append(max(ks) if ks else -1)
return out + [border_oracle(p)[-1]]
def test_example() -> None:
p = "abacabab"
assert compute_prefix_function(p) == [0, 0, 1, 0, 1, 2, 3, 2]
assert compute_next(p) == [-1, 0, -1, 1, -1, 0, -1, 3, 2]
assert knuth_morris_pratt("abacabcabacababa", p) == [7]
def test_edges() -> None:
assert knuth_morris_pratt("abc", "") == [0, 1, 2, 3]
assert knuth_morris_pratt("", "") == [0]
assert knuth_morris_pratt("", "a") == []
assert knuth_morris_pratt("ab", "abc") == []
assert knuth_morris_pratt("a", "a") == [0]
assert knuth_morris_pratt("a" * 10, "a" * 4) == list(range(7))
assert knuth_morris_pratt("a" * 10, "aaab") == []
def test_overlap_after_match() -> None:
# after a match the search must continue from nxt[m] = pi[m-1], not from 0
assert knuth_morris_pratt("abababab", "abab") == [0, 2, 4]
assert knuth_morris_pratt("aaaaa", "aa") == [0, 1, 2, 3]
def test_tables_against_definition() -> None:
rng = random.Random(77)
for alphabet in ("a", "ab", "abcd"):
for _ in range(300):
p = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 12)))
assert compute_prefix_function(p) == border_oracle(p)
assert compute_next(p) == next_oracle(p)
def test_random_against_oracle() -> None:
rng = random.Random(9001)
for alphabet in ("a", "ab", "abcd"):
for _ in range(700):
t = "".join(rng.choice(alphabet) for _ in range(rng.randint(0, 40)))
p = "".join(rng.choice(alphabet) for _ in range(rng.randint(0, 7)))
assert knuth_morris_pratt(t, p) == oracle(t, p)
def test_lists() -> None:
assert knuth_morris_pratt([0, 1, 0, 1, 0], [0, 1, 0]) == [0, 2]
What the next level needs
- the tests did not run in this build