Problem p02566 (python)

Fixed submission:
# Date [ 2020-09-09 16:51:26 ]
# Problem [ i.py ]
# Author Koki_tkg

import sys
# import math
# import bisect
# import numpy as np
# from decimal import Decimal
# from numba import njit, i8, u1, b1 #JIT compiler
# from itertools import combinations, product
# from collections import Counter, deque, defaultdict

# sys.setrecursionlimit(10 ** 6)
MOD = 10 ** 9 + 7
INF = 10 ** 9
PI = 3.14159265358979323846

def read_str():      return sys.stdin.readline().strip()
def read_int():      return int(sys.stdin.readline().strip())
def read_ints():     return map(int, sys.stdin.readline().strip().split())
def read_ints2(x):   return map(lambda num: int(num) - x, sys.stdin.readline().strip().split())
def read_str_list(): return list(sys.stdin.readline().strip().split())
def read_int_list(): return list(map(int, sys.stdin.readline().strip().split()))
def GCD(a: int, b: int) -> int: return b if a%b==0 else GCD(b, a%b)
def LCM(a: int, b: int) -> int: return (a * b) // GCD(a, b)

# Converted AC Library from C++ to python
def sa_native(s: list):
    from functools import cmp_to_key
    def mycmp(r, l):
        if l == r: return -1
        while l < n and r < n:
            if s[l] != s[r]: return 1 if s[l] < s[r] else -1
            l += 1
            r += 1
        return 1 if l == n else -1
    n = len(s)
    sa = [i for i in range(n)]
    sa.sort(key=cmp_to_key(mycmp))
    return sa

def sa_doubling(s: list):
    from functools import cmp_to_key
    def mycmp(y, x):
        if rnk[x] != rnk[y]: return 1 if rnk[x] < rnk[y] else -1
        rx = rnk[x + k] if x + k < n else -1
        ry = rnk[y + k] if y + k < n else -1
        return 1 if rx < ry else -1
        
    n = len(s)
    sa = [i for i in range(n)]; rnk = s; tmp = [0] * n; k = 1
    while k < n:
        sa.sort(key=cmp_to_key(mycmp))
        tmp[sa[0]] = 0
        for i in range(1, n):
            tmp[sa[i]] = tmp[sa[i - 1]]
            if mycmp(sa[i], sa[i - 1]) == 1: tmp[sa[i]] += 1
        tmp, rnk = rnk, tmp
        k *= 2
    return sa

def sa_is(s: list, upper: int):
    '''
    SA-IS, linear-time suffix array construction
    Reference:
    G. Nong, S. Zhang, and W. H. Chan, Two Efficient Algorithms for Linear Time Suffix Array Construction
    '''
    THRESHOLD_NATIVE = 10
    THRESHOLD_DOUBLING = 40
    n = len(s)
    if n == 0: return []
    if n == 1: return [0]
    if n == 2:
        if s[0] < s[1]:
            return [0, 1]
        else:
            return [1, 0]
    if n < THRESHOLD_NATIVE:
        return sa_native(s)
    if n < THRESHOLD_DOUBLING:
        return sa_doubling(s)
    
    sa = [0] * n
    ls = [False] * n
    for i in range(n - 2, -1, -1):
        ls[i] = ls[i + 1] if s[i] == s[i + 1] else s[i] < s[i + 1]
    
    sum_l = [0] * (upper + 1); sum_s = [0] * (upper + 1)
    for i in range(n):
        if not ls[i]:
            sum_s[s[i]] += 1
        else:
            sum_l[s[i] + 1] += 1
    for i in range(upper + 1):
        sum_s[i] += sum_l[i]
        if i < upper: sum_l[i + 1] += sum_s[i]
        
    def induce(lms: list):
        from copy import copy # ループより速く感じる
        for i in range(n): sa[i] = -1 # [-1]*n にするとバグる
        buf = [0] * (upper + 1)
        buf = copy(sum_s)
        for d in lms:
            if d == n: continue
            sa[buf[s[d]]] = d; buf[s[d]] += 1
        
        buf = copy(sum_l)
        sa[buf[s[n - 1]]] = n - 1; buf[s[n - 1]] += 1
        for i in range(n):
            v = sa[i]
            if v >= 1 and not ls[v - 1]:
                sa[buf[s[v - 1]]] = v - 1; buf[s[v - 1]] += 1
        
        buf = copy(sum_l)
        for i in range(n - 1, -1, -1):
            v = sa[i]
            if v >= 1 and ls[v - 1]:
                buf[s[v - 1] + 1] -= 1; sa[buf[s[v - 1] + 1]] = v - 1

    lms_map = [-1] * (n + 1)
    m = 0
    for i in range(1, n):
        if not ls[i - 1] and ls[i]:
            lms_map[i] = m; m += 1
    lms = [i for i in range(1, n) if not ls[i - 1] and ls[i]]
    induce(lms)

    if m:
        sorted_lms = [v for v in sa if lms_map[v] != -1]
        rec_s = [0] * m
        rec_upper = 0
        rec_s[lms_map[sorted_lms[0]]] = 0
        for i in range(1, m):
            l = sorted_lms[i - 1]; r = sorted_lms[i]
            end_l = lms[lms_map[l] + 1] if lms_map[l] + 1 < m else n
            end_r = lms[lms_map[r] + 1] if lms_map[r] + 1 < m else n
            same = True
            if end_l - l != end_r - r:
                same = False
            else:
                while l < end_l:
                    if s[l] != s[r]: break
                    l += 1
                    r += 1
                if l == n or s[l] != s[r]: same = False
            if not same: rec_upper += 1
            rec_s[lms_map[sorted_lms[i]]] = rec_upper
        
        rec_sa = sa_is(rec_s, rec_upper)

        for i in range(m):
            sorted_lms[i] = lms[rec_sa[i]]
        induce(sorted_lms)
    return sa

def suffix_array(s: list, upper: int):
    assert 0 <= upper
    for d in s:
        assert 0 <= d and d <= upper
    sa = sa_is(s, upper)
    return sa

def suffix_array2(s: list):
    from functools import cmp_to_key
    n = len(s)
    idx = [i for i in range(n)]
    idx.sort(key=cmp_to_key(lambda l, r: s[l] < s[r]))
    s2 = [0] * n
    now = 0
    for i in range(n):
        if i and s[idx[i - 1]] != s[idx[i]]: now += 1
        s2[idx[i]] = now
    return sa_is(s2, now)

def suffix_array3(s: str):
    n = len(s)
    s2 = list(map(ord, s))
    return sa_is(s2, 255)

def lcp_array(s: list, sa: list):
    '''
    Reference:
    T. Kasai, G. Lee, H. Arimura, S. Arikawa, and K. Park,
    Linear-Time Longest-Common-Prefix Computation in Suffix Arrays and Its
    Applications
    '''
    n = len(s)
    assert n >= 1
    rnk = [0] * n
    for i in range(n):
        rnk[sa[i]] = i
    lcp = [0] * (n - 1)
    h = 0
    for i in range(n):
        if h > 0: h -= 1
        if rnk[i] == 0: continue
        j = sa[rnk[i] - 1]
        while j + h < n and i + h < n:
            if s[j + h] != s[i + h]: break
            h += 1
        lcp[rnk[i] - 1] = h
    return lcp

def lcp_array2(s: str, sa: list):
    n = len(s)
    s2 = list(map(ord, s))
    return lcp_array(s2, sa)

def z_algorithm(s: list):
    '''
    Reference:
    D. Gusfield,
    Algorithms on Strings, Trees, and Sequences: Computer Science and
    Computational Biology
    '''
    n = len(s)
    if n == 0: return []
    z = [-1] * n
    z[0] = 0; j = 0
    for i in range(1, n):
        k = z[i] = 0 if j + z[j] <= i else min(j + z[j] - i, z[i - j])
        while i + k < n and s[k] == s[i + k]: k += 1
        z[i] = k
        if j + z[j] < i + z[i]: j = i
    z[0] = n
    return z

def z_algorithm2(s: str):
    n = len(s)
    s2 = list(map(ord, s))
    return z_algorithm(s2)

def Main():
    s = read_str()
    sa = suffix_array3(s)
    ans = len(s) * (len(s) + 1) // 2
    for x in lcp_array2(s, sa):
        ans -= x
    print(ans)

if __name__ == '__main__':
    Main()

Buggy submission (verdict: Runtime Error):
# Date [ 2020-09-09 19:51:26 ]
# Problem [ i.py ]
# Author Koki_tkg

import sys
# import math
# import bisect
# import numpy as np
# from decimal import Decimal
# from numba import njit, i8, u1, b1 #JIT compiler
# from itertools import combinations, product
# from collections import Counter, deque, defaultdict

# sys.setrecursionlimit(10 ** 6)
MOD = 10 ** 9 + 7
INF = 10 ** 9
PI = 3.14159265358979323846

def read_str():      return sys.stdin.readline().strip()
def read_int():      return int(sys.stdin.readline().strip())
def read_ints():     return map(int, sys.stdin.readline().strip().split())
def read_ints2(x):   return map(lambda num: int(num) - x, sys.stdin.readline().strip().split())
def read_str_list(): return list(sys.stdin.readline().strip().split())
def read_int_list(): return list(map(int, sys.stdin.readline().strip().split()))
def GCD(a: int, b: int) -> int: return b if a%b==0 else GCD(b, a%b)
def LCM(a: int, b: int) -> int: return (a * b) // GCD(a, b)

# Converted AC Library from C++ to python
def sa_native(s: list):
    from functools import cmp_to_key
    def mycmp(l, r):
        if l == r: return -1
        while l < n and r < n:
            if s[l] != s[r]: return 1 if s[l] < s[r] else -1
            l += 1
            r += 1
        return 1 if l == n else -1
    n = len(s)
    sa = [i for i in range(n)]
    sa.sort(key=cmp_to_key(mycmp))
    return sa

def sa_doubling(s: list):
    from functools import cmp_to_key
    def mycmp(x, y):
        if rnk[x] != rnk[y]: return 1 if rnk[x] < rnk[y] else -1
        rx = rnk[x + k] if x + k < n else -1
        ry = rnk[y + k] if y + k < n else -1
        return 1 if rx < ry else -1
        
    n = len(s)
    sa = [i for i in range(n)]; rnk = s; tmp = [0] * n; k = 1
    while k < n:
        sa.sort(key=cmp_to_key(mycmp))
        tmp[sa[0]] = 0
        for i in range(1, n):
            tmp[sa[i]] = tmp[sa[i - 1]]
            if mycmp(sa[i], sa[i - 1]) == 1: tmp[sa[i]] += 1
        tmp, rnk = rnk, tmp
        k *= 2
    return sa

def sa_is(s: list, upper: int):
    '''
    SA-IS, linear-time suffix array construction
    Reference:
    G. Nong, S. Zhang, and W. H. Chan, Two Efficient Algorithms for Linear Time Suffix Array Construction
    '''
    THRESHOLD_NATIVE = 10
    THRESHOLD_DOUBLING = 40
    n = len(s)
    if n == 0: return []
    if n == 1: return [0]
    if n == 2:
        if s[0] < s[1]:
            return [0, 1]
        else:
            return [1, 0]
    if n < THRESHOLD_NATIVE:
        return sa_native(s)
    if n < THRESHOLD_DOUBLING:
        return sa_doubling(s)
    
    sa = [0] * n
    ls = [False] * n
    for i in range(n - 2, -1, -1):
        ls[i] = ls[i + 1] if s[i] == s[i + 1] else s[i] < s[i + 1]
    
    sum_l = [0] * (upper + 1); sum_s = [0] * (upper + 1)
    for i in range(n):
        if not ls[i]:
            sum_s[s[i]] += 1
        else:
            sum_l[s[i] + 1] += 1
    for i in range(upper + 1):
        sum_s[i] += sum_l[i]
        if i < upper: sum_l[i + 1] += sum_s[i]
        
    def induce(lms: list):
        from copy import copy
        sa = [-1] * n
        buf = copy(sum_s)
        for d in lms:
            if d == n: continue
            sa[buf[s[d]]] = d; buf[s[d]] += 1
        
        buf = copy(sum_l)
        sa[buf[s[n - 1]]] = n - 1; buf[s[n - 1]] += 1
        for i in range(n):
            v = sa[i]
            if v >= 1 and not ls[v - 1]:
                sa[buf[s[v - 1]]] = v - 1; buf[s[v - 1]] += 1
        
        buf = copy(sum_l)
        for i in range(n - 1, -1, -1):
            v = sa[i]
            if v >= 1 and ls[v - 1]:
                buf[s[v - 1]] -= 1; sa[buf[s[v - 1]]] = v - 1
        buf.clear()

    lms_map = [-1] * (n + 1)
    m = 0
    for i in range(1, n):
        if not ls[i - 1] and ls[i]:
            lms_map[i] += m; m += 1
    
    lms = [i for i in range(1, n) if not ls[i - 1] and ls[i]]
    induce(lms)

    if m:
        sorted_lms = [v for v in sa if lms_map[v] != -1]
        rec_s = [-1] * m
        rec_upper = 0
        rec_s[lms_map[sorted_lms[0]]] = 0
        for i in range(1, m):
            l = sorted_lms[i - 1]; r = sorted_lms[i]
            end_l = lms[lms_map[l] + 1] if lms_map[l] + 1 < m else n
            end_r = lms[lms_map[r] + 1] if lms_map[r] + 1 < m else n
            same = True
            if end_l - l != end_r - r:
                same = False
            else:
                while l < end_l:
                    if s[l] != s[r]: break
                    l += 1
                    r += 1
                if l == n or s[l] != s[r]: same = False
            if not same: rec_upper += 1
            rec_s[lms_map[sorted_lms[i]]] = rec_upper
        
        rec_sa = sa_is(rec_s, rec_upper)

        for i in range(m):
            sorted_lms[i] = lms[rec_sa[i]]
        induce(sorted_lms)
    return sa

def suffix_array(s: list, upper: int):
    assert 0 <= upper
    for d in s:
        assert 0 <= d and d <= upper
    sa = sa_is(s, upper)
    return sa

def suffix_array2(s: list):
    from functools import cmp_to_key
    n = len(s)
    idx = [i for i in range(n)]
    idx.sort(key=cmp_to_key(lambda l, r: s[l] < s[r]))
    s2 = [0] * n
    now = 0
    for i in range(n):
        if i and s[idx[i - 1]] != s[idx[i]]: now += 1
        s2[idx[i]] = now
    return sa_is(s2, now)

def suffix_array3(s: str):
    n = len(s)
    s2 = list(map(ord, s))
    return sa_is(s2, 255)

def lcp_array(s: list, sa: list):
    '''
    Reference:
    T. Kasai, G. Lee, H. Arimura, S. Arikawa, and K. Park,
    Linear-Time Longest-Common-Prefix Computation in Suffix Arrays and Its
    Applications
    '''
    n = len(s)
    assert n >= 1
    rnk = [0] * n
    for i in range(n):
        rnk[sa[i]] = i
    lcp = [0] * (n - 1)
    h = 0
    for i in range(n):
        if h > 0: h -= 1
        if rnk[i] == 0: continue
        j = sa[rnk[i] - 1]
        while j + h < n and i + h < n:
            if s[j + h] != s[i + h]: break
            h += 1
        lcp[rnk[i] - 1] = h
    return lcp

def lcp_array2(s: str, sa: list):
    n = len(s)
    s2 = list(map(ord, s))
    return lcp_array(s2, sa)

def z_algorithm(s: list):
    '''
    Reference:
    D. Gusfield,
    Algorithms on Strings, Trees, and Sequences: Computer Science and
    Computational Biology
    '''
    n = len(s)
    if n == 0: return []
    z = [-1] * n
    z[0] = 0; j = 0
    for i in range(1, n):
        k = z[i] = 0 if j + z[j] <= i else min(j + z[j] - i, z[i - j])
        while i + k < n and s[k] == s[i + k]: k += 1
        z[i] = k
        if j + z[j] < i + z[i]: j = i
    z[0] = n
    return z

def z_algorithm2(s: str):
    n = len(s)
    s2 = list(map(ord, s))
    return z_algorithm(s2)

def Main():
    s = read_str()
    sa = suffix_array3(s)
    ans = len(s) * (len(s) + 1) // 2
    for x in lcp_array2(s, sa):
        ans -= x
    print(ans)

if __name__ == '__main__':
    Main()