Skip to content

Commit 40e7bfd

Browse files
committed
feat(strings): add suffix array algorithm
1 parent 574d100 commit 40e7bfd

1 file changed

Lines changed: 70 additions & 0 deletions

File tree

‎strings/suffix_array.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
"""
2+
https://en.wikipedia.org/wiki/Suffix_array
3+
https://cp-algorithms.com/string/suffix-array.html
4+
5+
A suffix array lists the starting positions of every suffix of a string in
6+
lexicographic order. It needs one integer per suffix, where a suffix tree
7+
needs a node per character, so it is the compact structure used for pattern
8+
search and for the LCP-based algorithms built on top of it.
9+
10+
Construction here uses prefix doubling (Manber-Myers).
11+
12+
Time Complexity: O(n * log(n)^2) - where n is the length of the text
13+
Space Complexity: O(n)
14+
"""
15+
16+
17+
def suffix_array(text: str) -> list[int]:
18+
"""
19+
Return the starting indices of all suffixes of `text`, sorted.
20+
21+
>>> suffix_array("banana")
22+
[5, 3, 1, 0, 4, 2]
23+
>>> suffix_array("mississippi")
24+
[10, 7, 4, 1, 0, 9, 8, 6, 3, 5, 2]
25+
>>> suffix_array("abababab")
26+
[6, 4, 2, 0, 7, 5, 3, 1]
27+
>>> suffix_array("aaaa")
28+
[3, 2, 1, 0]
29+
>>> suffix_array("aça")
30+
[2, 0, 1]
31+
>>> suffix_array("a")
32+
[0]
33+
>>> suffix_array("")
34+
[]
35+
"""
36+
n = len(text)
37+
if n == 0:
38+
return []
39+
40+
suffixes = list(range(n))
41+
rank = list(map(ord, text))
42+
k = 1
43+
44+
while k < n:
45+
keys = [
46+
(rank[i], rank[i + k] if i + k < n else -1, i)
47+
for i in range(n)
48+
]
49+
keys.sort()
50+
suffixes = [i for _, _, i in keys]
51+
52+
new_rank = [0] * n
53+
for position in range(1, n):
54+
previous, current = suffixes[position - 1], suffixes[position]
55+
previous_key = (
56+
rank[previous],
57+
rank[previous + k] if previous + k < n else -1,
58+
)
59+
current_key = (
60+
rank[current],
61+
rank[current + k] if current + k < n else -1,
62+
)
63+
new_rank[current] = new_rank[previous] + int(previous_key != current_key)
64+
65+
rank = new_rank
66+
if rank[suffixes[-1]] == n - 1:
67+
break
68+
k *= 2
69+
70+
return suffixes

0 commit comments

Comments
 (0)