beam: suffix automaton with endpos counts and LCS, tested

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())

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")