[Python] Python Suffix Array (O(N log N)) Causes Runtime Error on Codeforces (ITMO Pilot...

I'm implementing a suffix array using the "doubling" algorithm with radix sort, as taught in the ITMO University Pilot Course (Step 2 on Suffix Arrays). My Python code, which I believe has an O(NlogN) time complexity, passes small local tests but hits a "Runtime Error" on Codeforces. I've attached the full error message and my code.

Codeforces Problem: ITMO Academy: Suffix Array - Step 2

Complete Error Message from Codeforces:

→Judgement Protocol Test: #1, time: 77 ms., memory: 2792 KB, exit code: 1, checker exit code: 0, verdict: RUNTIME_ERROR Input ababba Checker Log Exit code is 1


Python

import itertools
import math


def next_power_of_2(n):
return 1 if n == 0 else 2 ** math.ceil(math.log2(n))


def custom_compare(str1, str2, char_rank):
for c1, c2 in zip(str1, str2):
if char_rank[c1] < char_rank[c2]:
return True
elif char_rank[c1] > char_rank[c2]:
return False
return len(str1) < len(str2)


def counting_sort_by_key(arr, key_func, max_val):
n = len(arr)
count = [0] * (max_val + 1)
output = [0] * n

for item in arr:
count[key_func(item)] += 1

for i in range(1, len(count)):
count += count[i-1]

for item in reversed(arr):
count[key_func(item)] -= 1
output[count[key_func(item)]] = item

return output


def main():
base = input() + '$'
n = len(base)

# PROBLEM LIKELY HERE: char_rank is fixed, but alphabet might contain other characters
char_rank = {'$': 0, 'a': 1, 'b': 2}

# --- The following block appears to be an attempt at a different suffix array approach or
# preparation that might not be necessary for the doubling algorithm and is potentially inefficient.
# The core doubling algorithm starts further down.
suffixes = [base[i:] for i in range(len(base))]
max_len = next_power_of_2(max(len(suffix) for suffix in suffixes))

for i in range(len(suffixes)):
current = suffixes
needed = max_len - len(current)
if needed > 0:
try:
idx = base.index(current)
cycle_start = base[idx + len(current):] + base[:idx + len(current)]
except ValueError:
cycle_start = base
extension = ''.join(itertools.islice(itertools.cycle(cycle_start), needed))
suffixes += extension

arr = [int(math.pow(2, i)) for i in range(int(math.log2(max_len)) + 1)]

df = []
ansi = []
for j in arr:
ds = [suffix[:j] for suffix in suffixes]
ds = [[s[:j//2], s[j//2:]] for s in ds]
df.append(ds)
comparisons = []
for k in range(len(ds) - 1):
str1 = ''.join(ds[k])
str2 = ''.join(ds[k + 1])
result = custom_compare(str1, str2, char_rank)
comparisons.append(result)
ansi.append(comparisons)
# --- End of potentially unnecessary/inefficient block

# Initial sort by single characters (k=0 step)
p = list(range(n))
# PROBLEM LIKELY HERE: If base is not in char_rank, it will raise a KeyError
c = [char_rank[base] for i in range(n)]
p.sort(key=lambda i: c)

new_c = [0] * n
rank = 0
new_c[p[0]] = 0
for i in range(1, n):
if c[p] != c[p[i-1]]:
rank += 1
new_c[p] = rank
c = new_c

# Doubling steps
for k in range(int(math.log2(max_len))): # Should be math.ceil(math.log2(n)) or similar for correct number of steps
shift = 1 << k
# Sort by (c, c[(i + shift) % n]) using counting sort
pairs = [(c, c[(i + shift) % n], i) for i in range(n)]

max_val = max(c) # Max value for the first component for counting sort

# This sorts by the first component (c)
pairs = counting_sort_by_key(pairs, lambda x: x[0], max_val)

p = [pair[2] for pair in pairs] # p now contains indices sorted by the first component

new_c = [0] * n
rank = 0
new_c[p[0]] = 0
for i in range(1, n):
# This comparison is for the new equivalence classes based on the sorted pairs
if pairs[:2] != pairs[i-1][:2]:
rank += 1
new_c[p] = rank
c = new_c

# This comparison block at the end of the loop seems extraneous to the core suffix array construction.
# It's calculating custom_compare results but not using them to update 'p' or 'c'.
for i in range(n-1):
str1 = base[p+2**(k+1)] if p+2**(k+1) <= n else base[p:] + base[:2**(k+1)-n+p]
str2 = base[p[i+1][i+1]+2**(k+1)] if p[i+1]+2**(k+1) <= n else base[p[i+1]:] + base[:2**(k+1)-n+p[i+1]]
result = custom_compare(str1, str2, char_rank)

print(' '.join(map(str, p)))


if __name__ == "__main__":
main()


My Implementation Details:

I'm building the suffix array iteratively by doubling the length of the considered substrings. For sorting, I'm using a counting_sort_by_key function, which acts as a radix sort implementation.

Specifically, the counting_sort_by_key function sorts elements based on a given key. In the main loop, I'm trying to sort pairs (c, c[(i + shift) % n], i). My current approach calls counting_sort_by_key using lambda x: x[0] to sort by the first component of these pairs.

The char_rank dictionary is currently initialized with fixed values for '$', 'a', and 'b'.

What I've tried:

Verified the counting_sort_by_key implementation separately. Checked the logic for handling cyclic shifts and padding with $ or extensions for suffixes (though I suspect this section might be unnecessary for the doubling algorithm). Reviewed the doubling logic and how c (equivalence classes/ranks) is updated.

Just to clear confusion:

i know that char_rank is fixed but my the current testcase only contain ababba which is a and b and is present in char_rank so problem must be in some other part of code

Continue reading...