From 5f3bfecb264875c06cfe551843e92027bb7ba466 Mon Sep 17 00:00:00 2001 From: beam Date: Sun, 1 Nov 2026 08:48:44 +0000 Subject: [PATCH] beam: suffix automaton with endpos counts and LCS, tested --- wizard-work/beam/suffix_automaton.py | 121 ++++++++++++++++++++++ wizard-work/beam/test_suffix_automaton.py | 52 ++++++++++ 2 files changed, 173 insertions(+) create mode 100644 wizard-work/beam/suffix_automaton.py create mode 100644 wizard-work/beam/test_suffix_automaton.py diff --git a/wizard-work/beam/suffix_automaton.py b/wizard-work/beam/suffix_automaton.py new file mode 100644 index 0000000..4971ac5 --- /dev/null +++ b/wizard-work/beam/suffix_automaton.py @@ -0,0 +1,121 @@ +"""Suffix Automaton (SAM) with endpos occurrence counts and distinct-substring counting. + +A suffix automaton for a string s is a minimal DFA recognizing all suffixes of s +(equivalently all substrings). Each state corresponds to an equivalence class of +substrings sharing the same set of end positions (endpos set). + +Correctness notes (the subtle parts): +- On extension, the new state `cur` represents the whole prefix s[0..i]; it gets + cnt=1 (it is an end position). Cloned states get cnt=0 because they are not + new end positions, only split classes. +- The clone copies transitions/len/link from q, with len = p.len + 1. +- After all extensions, occurrence counts are propagated in decreasing order of + `len` (topological order of the suffix links): a state's count flows to its + suffix link. This yields, for every state, the size of its endpos set. +- Distinct substrings = sum over states of (state.len - state.link.len), since + each state contributes the substrings whose lengths lie in (link.len, len]. +- Longest common substring of s and t: walk t through the automaton built from s, + tracking current state and current match length; reset on missing transition. +""" + +from collections import defaultdict, deque + + +class SuffixAutomaton: + def __init__(self, s: str): + # states: list of dicts + self.states = [{"len": 0, "link": -1, "next": {}, "cnt": 0}] + self.last = 0 + for ch in s: + self._extend(ch) + self._compute_cnt() + self.s = s + + def _extend(self, ch): + st = self.states + cur = len(st) + st.append({"len": st[self.last]["len"] + 1, "link": 0, "next": {}, "cnt": 1}) + p = self.last + while p != -1 and ch not in st[p]["next"]: + st[p]["next"][ch] = cur + p = st[p]["link"] + if p == -1: + st[cur]["link"] = 0 + else: + q = st[p]["next"][ch] + if st[p]["len"] + 1 == st[q]["len"]: + st[cur]["link"] = q + else: + clone = len(st) + st.append({ + "len": st[p]["len"] + 1, + "link": st[q]["link"], + "next": dict(st[q]["next"]), + "cnt": 0, # clone is NOT a new end position + }) + while p != -1 and st[p]["next"].get(ch) == q: + st[p]["next"][ch] = clone + p = st[p]["link"] + st[q]["link"] = clone + st[cur]["link"] = clone + self.last = cur + + def _compute_cnt(self): + st = self.states + # bucket sort by len + maxlen = max(x["len"] for x in st) + buckets = [[] for _ in range(maxlen + 1)] + for i, x in enumerate(st): + buckets[x["len"]].append(i) + # process lengths descending + for L in range(maxlen, 0, -1): + for v in buckets[L]: + link = st[v]["link"] + if link != -1 and link != v: + st[link]["cnt"] += st[v]["cnt"] + + def distinct_substrings(self) -> int: + total = 0 + for v in range(1, len(self.states)): + total += self.states[v]["len"] - self.states[self.states[v]["link"]]["len"] + return total + + def occurrences(self, sub: str) -> int: + # number of (possibly overlapping) occurrences of sub in s + v = 0 + for ch in sub: + if ch not in self.states[v]["next"]: + return 0 + v = self.states[v]["next"][ch] + return self.states[v]["cnt"] + + def longest_common_substring(self, t: str) -> str: + v = 0 + l = 0 + best_len = 0 + best_end = 0 # end index in t (exclusive) + for i, ch in enumerate(t): + while v != 0 and ch not in self.states[v]["next"]: + v = self.states[v]["link"] + l = self.states[v]["len"] + if ch in self.states[v]["next"]: + v = self.states[v]["next"][ch] + l += 1 + else: + v = 0 + l = 0 + if l > best_len: + best_len = l + best_end = i + 1 + return t[best_end - best_len: best_end] + + +def build(s): + return SuffixAutomaton(s) + + +if __name__ == "__main__": + import sys + s = sys.argv[1] if len(sys.argv) > 1 else "ababa" + sam = SuffixAutomaton(s) + print("distinct:", sam.distinct_substrings()) diff --git a/wizard-work/beam/test_suffix_automaton.py b/wizard-work/beam/test_suffix_automaton.py new file mode 100644 index 0000000..cf8bdf6 --- /dev/null +++ b/wizard-work/beam/test_suffix_automaton.py @@ -0,0 +1,52 @@ +import itertools +from suffix_automaton import SuffixAutomaton + +def brute_distinct(s): + subs = set() + for i in range(len(s)): + for j in range(i+1, len(s)+1): + subs.add(s[i:j]) + return len(subs) + +def brute_occ(s, sub): + if not sub: + return len(s)+1 + c = 0 + i = 0 + while True: + i = s.find(sub, i) + if i == -1: + break + c += 1 + i += 1 + return c + +def test_distinct(): + for s in ["", "a", "ab", "aba", "ababa", "aaaa", "abcabc", "mississippi"]: + sam = SuffixAutomaton(s) + assert sam.distinct_substrings() == brute_distinct(s), (s, sam.distinct_substrings(), brute_distinct(s)) + +def test_occurrences(): + for s in ["ababa", "aaaa", "abcabc", "mississippi", "banana"]: + sam = SuffixAutomaton(s) + for L in range(1, 4): + for sub in set(s[i:i+L] for i in range(len(s))): + assert sam.occurrences(sub) == brute_occ(s, sub), (s, sub) + +def test_lcs(): + cases = [("abcd", "bcde", "bcd"), ("abc", "def", ""), ("aaaa", "aa", "aa"), + ("ababc", "bab", "bab"), ("", "abc", "")] + for a, b, exp in cases: + sam = SuffixAutomaton(a) + got = sam.longest_common_substring(b) + assert got == exp, (a, b, got, exp) + +def test_empty(): + sam = SuffixAutomaton("") + assert sam.distinct_substrings() == 0 + assert sam.occurrences("x") == 0 + assert sam.longest_common_substring("abc") == "" + +if __name__ == "__main__": + test_distinct(); test_occurrences(); test_lcs(); test_empty() + print("ALL PASS")