# Source code for tryalgo.suffix_array

```#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""\
suffix array,
but only the O(n log^2(n)) implementation, which is enough for most programming contest problems

christoph dürr 2024
"""

[docs]
def sort_class(s):
""" sorts s and returns additional information

:param s: string or list
:returns p, c: p[j]=i if s[i] has rank j in sorted(s) and c[i] is rank of s[i] in sorted(set(s))
:complexity: O(n log n) or better if sort makes use of specific values in s
"""
S_index = [(x, i) for i, x in enumerate(s)]
p = [i for x, i in sorted(S_index)]
x2c = {x : i for i, x in enumerate(sorted(set(s)))}
c = [x2c[x] for x in s]
return p, c

[docs]
def sort_cyclic_shifts(s):
""" given a string s, sort lexicographically all cyclic shifts of s.

The i-th cyclic shift of s is s[i:] + s[i:]
:param s: string or list
:returns L: such that L[j]=i if the i-th cyclic shift of s has rank j
:complexity: O(n * log(n)^2)
"""
p, c = sort_class(s)
n = len(s)
K = 1
while K <= n:
L = [(c[i], c[(i + K) % n]) for i in range(n)]
p, c = sort_class(L)
K <<= 1
return p

[docs]
def suffix_array(s):
""" given a string s, sort lexicographically suffixes of s
:param s: string
:returns: R with R[i] is j such that s[j:] has rank i
:complexity: O(n log^2 n)
"""
special = chr(0)
assert special < min(s)
L = sort_cyclic_shifts(s + special)
return L[1:]

if __name__ == "__main__":
# tested at https://www.spoj.com/problems/SARRAY/
import sys