Suffix array pattern search in Python
Finds a pattern by binary search over the suffix array, in O(m log n) time, or O(m + log n) with LCP values.
- File
suffix_array_search.py, 99 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/suffix-array-search/. The authors are listed in prov/VMB6DY.jsonl.
from collections.abc import Sequence
from typing import Any
def lcp_lr(lcp: Sequence[int]) -> tuple[list[int], list[int]]:
"""The values Llcp[M] and Rlcp[M] of the fixed binary search on (L, R).
lcp[r] is the longest common prefix of the suffixes of ranks r - 1 and r
(lcp[0] is unused). The search starts with L = -1 and R = n, M is
(L + R) // 2, and the virtual ranks -1 and n have lcp 0 with everything.
Llcp[M] = lcp of ranks L and M, Rlcp[M] = lcp of ranks M and R. Time O(n),
with an explicit stack instead of recursion.
"""
n = len(lcp)
llcp, rlcp = [0] * n, [0] * n
def edge(a: int, b: int) -> int: # lcp of two neighbor ranks, 0 for a virtual one
return lcp[b] if 0 <= a and b < n else 0
result: dict[tuple[int, int], int] = {}
stack = [(-1, n, False)]
while stack:
lo, hi, done = stack.pop()
if hi - lo == 1:
result[(lo, hi)] = edge(lo, hi)
continue
mid = (lo + hi) // 2
if not done:
stack.append((lo, hi, True))
stack.append((lo, mid, False))
stack.append((mid, hi, False))
continue
a = result.pop((lo, mid))
b = result.pop((mid, hi))
llcp[mid] = a if lo >= 0 else 0
rlcp[mid] = b if hi < n else 0
result[(lo, hi)] = min(a, b) if lo >= 0 and hi < n else 0
return llcp, rlcp
def sa_bound(s: Sequence[Any], sa: Sequence[int], llcp: Sequence[int], rlcp: Sequence[int],
p: Sequence[Any], strict: bool, counter: list[int] | None = None) -> int:
"""The first rank whose suffix has its len(p)-prefix >= p (or > p if strict).
Symbols are compared with == and < only (a general ordered alphabet).
Positions and ranks are 0-based; no sentinel. At most
len(p) + 2 ceil(log2(n + 1)) symbol comparisons. If counter is given,
counter[0] is increased by the number of symbol comparisons.
"""
n, m = len(sa), len(p)
lo, hi = -1, n # @line L1: the suffix of rank lo is < p, the one of rank hi is >= p
l = r = 0 # @line L2: l = lcp(p, suffix lo), r = lcp(p, suffix hi)
while hi - lo > 1: # @line L3
mid = (lo + hi) // 2 # @line L4
if l >= r: # @line L5
if llcp[mid] > l: # @line L6
lo = mid # the suffix of mid is like the one of lo
continue
if llcp[mid] < l: # @line L7
hi, r = mid, llcp[mid]
continue
k = l # @line L8
else:
if rlcp[mid] > r: # @line L9
hi = mid
continue
if rlcp[mid] < r: # @line L10
lo, l = mid, rlcp[mid]
continue
k = r # @line L11
pos = sa[mid]
while k < m and pos + k < n: # @line L12: compare from the known lcp on
if counter is not None:
counter[0] += 1
if p[k] != s[pos + k]:
break
k += 1
if k == m: # @line L13: the suffix starts with p
go_right = strict
elif pos + k == n: # the suffix is a proper prefix of p
go_right = True
else:
if counter is not None:
counter[0] += 1
go_right = s[pos + k] < p[k]
if go_right: # @line L14
lo, l = mid, k
else:
hi, r = mid, k # @line L15
return hi # @line L16
def sa_range(s: Sequence[Any], sa: Sequence[int], lcp: Sequence[int], p: Sequence[Any]) -> tuple[int, int]:
"""The half-open rank interval of the suffixes that start with p."""
llcp, rlcp = lcp_lr(lcp)
return sa_bound(s, sa, llcp, rlcp, p, False), sa_bound(s, sa, llcp, rlcp, p, True)
test_suffix_array_search.py 113 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/suffix-array-search/. The authors are listed in prov/VMB6DY.jsonl.
import math
import random
from suffix_array_search import lcp_lr, sa_bound, sa_range
def naive_sa(s: str) -> list[int]:
return sorted(range(len(s)), key=lambda i: s[i:])
def naive_lcp(s: str, sa: list[int]) -> list[int]:
out = [0] * len(sa)
for r in range(1, len(sa)):
a, b, k = s[sa[r - 1]:], s[sa[r]:], 0
while k < len(a) and k < len(b) and a[k] == b[k]:
k += 1
out[r] = k
return out
def naive_range(s: str, sa: list[int], p: str) -> tuple[int, int]:
heads = [s[i:i + len(p)] for i in sa]
lo = sum(1 for h in heads if h < p)
return lo, lo + sum(1 for h in heads if h == p)
def setup(s: str) -> tuple[list[int], list[int]]:
sa = naive_sa(s)
return sa, naive_lcp(s, sa)
def test_example() -> None:
s = "mississippi"
sa, lcp = setup(s)
assert sa == [10, 7, 4, 1, 0, 9, 8, 6, 3, 5, 2]
assert lcp == [0, 1, 1, 4, 0, 0, 1, 0, 2, 1, 3]
assert sa_range(s, sa, lcp, "ssi") == (9, 11)
assert sa_range(s, sa, lcp, "issi") == (2, 4)
assert sa_range(s, sa, lcp, "i") == (0, 4)
assert sa_range(s, sa, lcp, "sp") == (9, 9)
assert sa_range(s, sa, lcp, "z") == (11, 11)
assert sa_range(s, sa, lcp, "a") == (0, 0)
llcp, rlcp = lcp_lr(lcp)
assert llcp[5] == 0 and rlcp[5] == 0
assert llcp[8] == 0 and rlcp[8] == 0
assert llcp[6] == 1 and rlcp[6] == 0
def test_lcp_lr_against_naive() -> None:
rng = random.Random(5)
for _ in range(300):
n = rng.randint(0, 40)
s = "".join(rng.choice("ab"[: rng.choice([1, 2])]) for _ in range(n))
sa, lcp = setup(s)
llcp, rlcp = lcp_lr(lcp)
def pair(a: int, b: int) -> int:
if a < 0 or b >= n:
return 0
return min(lcp[a + 1:b + 1])
stack = [(-1, n)]
while stack:
lo, hi = stack.pop()
if hi - lo <= 1:
continue
mid = (lo + hi) // 2
assert llcp[mid] == pair(lo, mid) and rlcp[mid] == pair(mid, hi)
stack += [(lo, mid), (mid, hi)]
def test_edges() -> None:
assert sa_range("", [], [], "a") == (0, 0)
assert sa_range("a", [0], [0], "a") == (0, 1)
assert sa_range("a", [0], [0], "aa") == (1, 1)
s = "aaaaaaaa"
sa, lcp = setup(s)
assert sa_range(s, sa, lcp, "aaa") == (2, 8)
assert sa_range(s, sa, lcp, "a" * 9) == (8, 8)
def test_comparison_bound() -> None:
rng = random.Random(77)
for s in ("a" * 400, "ab" * 200, "".join(rng.choice("ab") for _ in range(400))):
sa, lcp = setup(s)
llcp, rlcp = lcp_lr(lcp)
for m in (1, 5, 50, 150):
i = rng.randrange(len(s) - m)
for p in (s[i:i + m], s[i:i + m - 1] + "c", "a" * m):
for strict in (False, True):
c = [0]
sa_bound(s, sa, llcp, rlcp, p, strict, c)
assert c[0] <= m + 2 * math.ceil(math.log2(len(s) + 1))
def test_random_against_naive() -> None:
rng = random.Random(1993)
for trial in range(700):
sigma = "abcd"[: rng.choice([1, 2, 4])]
n = rng.randint(0, 30)
s = "".join(rng.choice(sigma) for _ in range(n))
if trial % 35 == 0:
s = rng.choice(["a" * n, "ab" * (n // 2), ""])
sa, lcp = setup(s)
for _ in range(10):
p = "".join(rng.choice(sigma + "e") for _ in range(rng.randint(1, 6)))
if s and rng.random() < 0.6:
i = rng.randrange(len(s))
p = s[i:rng.randint(i + 1, len(s))]
assert sa_range(s, sa, lcp, p) == naive_range(s, sa, p)
What the next level needs
- the tests did not run in this build