122 lines
4.3 KiB
Python
122 lines
4.3 KiB
Python
"""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())
|