53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
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")
|