Problem:


Score : 100 points

Problem Statement

You are given a string of length N. Calculate the number of distinct substrings of S.

Constraints

1 \leq N \leq 500,000

S consists of lowercase English letters.

Input

Input is given from Standard Input in the following format:

S

Output

Print the answer.

Sample Input 1

abcbcba

Sample Output 1

21

Sample Input 2

mississippi

Sample Output 2

53

Sample Input 3

ababacaca

Sample Output 3

33

Sample Input 4

aaaaa

Sample Output 4

5



Sample input:
abcbcba


Sample output:
21

Submission status: Accepted

Solution code:
import types

_atcoder_code = """
# Python port of AtCoder Library.

__version__ = '0.0.1'
"""

atcoder = types.ModuleType('atcoder')
exec(_atcoder_code, atcoder.__dict__)

_atcoder_string_code = """
import copy
import functools
import typing


def _sa_naive(s: typing.List[int]) -> typing.List[int]:
    sa = list(range(len(s)))
    return sorted(sa, key=lambda i: s[i:])


def _sa_doubling(s: typing.List[int]) -> typing.List[int]:
    n = len(s)
    sa = list(range(n))
    rnk = copy.deepcopy(s)
    tmp = [0] * n
    k = 1
    while k < n:
        def cmp(x: int, y: int) -> bool:
            if rnk[x] != rnk[y]:
                return rnk[x] - rnk[y]
            rx = rnk[x + k] if x + k < n else -1
            ry = rnk[y + k] if y + k < n else -1
            return rx - ry
        sa.sort(key=functools.cmp_to_key(cmp))
        tmp[sa[0]] = 0
        for i in range(1, n):
            tmp[sa[i]] = tmp[sa[i - 1]] + (1 if cmp(sa[i - 1], sa[i]) else 0)
        tmp, rnk = rnk, tmp
        k *= 2
    return sa


def _sa_is(s: typing.List[int], upper: int) -> typing.List[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_naive = 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_naive:
        return _sa_naive(s)
    if n < threshold_doubling:
        return _sa_doubling(s)

    sa = [0] * n
    ls = [False] * n
    for i in range(n - 2, -1, -1):
        if s[i] == s[i + 1]:
            ls[i] = ls[i + 1]
        else:
            ls[i] = 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: typing.List[int]) -> None:
        nonlocal sa
        sa = [-1] * n

        buf = copy.deepcopy(sum_s)
        for d in lms:
            if d == n:
                continue
            sa[buf[s[d]]] = d
            buf[s[d]] += 1

        buf = copy.deepcopy(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.deepcopy(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 = []
    for i in range(1, n):
        if not ls[i - 1] and ls[i]:
            lms.append(i)

    induce(lms)

    if m:
        sorted_lms = []
        for v in sa:
            if lms_map[v] != -1:
                sorted_lms.append(v)
        rec_s = [0] * m
        rec_upper = 0
        rec_s[lms_map[sorted_lms[0]]] = 0
        for i in range(1, m):
            left = sorted_lms[i - 1]
            right = sorted_lms[i]
            if lms_map[left] + 1 < m:
                end_l = lms[lms_map[left] + 1]
            else:
                end_l = n
            if lms_map[right] + 1 < m:
                end_r = lms[lms_map[right] + 1]
            else:
                end_r = n

            same = True
            if end_l - left != end_r - right:
                same = False
            else:
                while left < end_l:
                    if s[left] != s[right]:
                        break
                    left += 1
                    right += 1
                if left == n or s[left] != s[right]:
                    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: typing.Union[str, typing.List[int]],
                 upper: typing.Optional[int] = None) -> typing.List[int]:
    if isinstance(s, str):
        return _sa_is([ord(c) for c in s], 255)
    elif upper is None:
        n = len(s)
        idx = list(range(n))
        idx.sort(key=functools.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)
    else:
        assert 0 <= upper
        for d in s:
            assert 0 <= d <= upper

        return _sa_is(s, upper)


def lcp_array(s: typing.Union[str, typing.List[int]],
              sa: typing.List[int]) -> typing.List[int]:
    '''
    Reference:
    T. Kasai, G. Lee, H. Arimura, S. Arikawa, and K. Park,
    Linear-Time Longest-Common-Prefix Computation in Suffix Arrays and Its
    Applications
    '''

    if isinstance(s, str):
        s = [ord(c) for c in s]

    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 z_algorithm(s: typing.Union[str, typing.List[int]]) -> typing.List[int]:
    '''
    Reference:
    D. Gusfield,
    Algorithms on Strings, Trees, and Sequences: Computer Science and
    Computational Biology
    '''

    if isinstance(s, str):
        s = [ord(c) for c in s]

    n = len(s)
    if n == 0:
        return []

    z = [0] * n
    j = 0
    for i in range(1, n):
        z[i] = 0 if j + z[j] <= i else min(j + z[j] - i, z[i - j])
        while i + z[i] < n and s[z[i]] == s[i + z[i]]:
            z[i] += 1
        if j + z[j] < i + z[i]:
            j = i
    z[0] = n

    return z
"""

atcoder.string = types.ModuleType('atcoder.string')
exec(_atcoder_string_code, atcoder.string.__dict__)
suffix_array = atcoder.string.suffix_array

lcp_array = atcoder.string.lcp_array

# https://atcoder.jp/contests/practice2/tasks/practice2_i

import sys

# from atcoder.string import suffix_array, lcp_array


def main() -> None:
    s = sys.stdin.readline().strip()
    sa = suffix_array(s)

    answer = len(s) * (len(s) + 1) // 2
    for x in lcp_array(s, sa):
        answer -= x
    print(answer)


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
mod = 1000000007
eps = 10**-9


def main():
    import sys
    from itertools import accumulate
    input = sys.stdin.readline

    # O(NlogN)
    # return suffix array and lcp array
    # 0th index of SA is |S|
    def SuffixArray(S):
        S = S + "$"
        N = len(S)
        LV = N.bit_length()

        # 0th sort
        P = []
        C = [0] * N
        cnt = [[] for _ in range(256)]
        for i, s in enumerate(S):
            cnt[ord(s)].append(i)
        cnt_c = 0
        for i_list in cnt:
            for i in i_list:
                P.append(i)
                C[i] = cnt_c
            if i_list:
                cnt_c += 1

        # doubling
        for k in range(LV):
            C_new = [0] * N
            cnt = [0] * N
            k2 = 1 << k
            for i in range(N):
                p = P[i] - k2
                if p < 0:
                    p += N
                cnt[C[p]] += 1
                P[i] = p
            cnt = list(accumulate(cnt))
            P_new = [0] * N
            for i in range(N-1, -1, -1):
                p = P[i]
                c = C[p]
                P_new[cnt[c] - 1] = p
                cnt[c] -= 1
            cnt_c = -1
            prev = None
            for p in P_new:
                p_ = p + k2
                if p_ >= N:
                    p_ -= N
                new = C[p] * N + C[p_]
                if new != prev:
                    cnt_c += 1
                    prev = new
                C_new[p] = cnt_c
            C = C_new
            P = P_new

        # lcp
        rank = [0] * N
        for i in range(N):
            rank[P[i]] = i
        LCP = [0] * (N - 1)
        h = 0
        for i in range(N - 1):
            j = P[rank[i] - 1]
            if h > 0:
                h -= 1
            while i + h < N and j + h < N:
                if S[i + h] == S[j + h]:
                    h += 1
                else:
                    break
            LCP[rank[i] - 1] = h

        return P, LCP

    S = input().rstrip('\n')
    N = len(S)
    SA, LCP = SuffixArray(S)
    print(N * (N+1) // 2 - sum(LCP))


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline

def SA_IS(s, upper=127):
    n = len(s)
    sa = [-1 for _ in range(n)]
    bucket = [0 for _ in range(upper + 1)]
    smaller = [True for _ in range(n)]
    is_lms = [False for _ in range(n)]
    lmss = []

    for x in s:
        bucket[x+1] += 1
    for i in range(upper):
        bucket[i+1] += bucket[i]
    
    for i in range(n-2, -1, -1):
        if s[i] > s[i+1] or (s[i] == s[i+1] and not smaller[i+1]):
            smaller[i] = False
            if smaller[i+1]:
                is_lms[i+1] = True
                lmss.append(i+1)

    count = [0 for _ in range(upper + 1)]
    for i in lmss:
        sa[bucket[s[i] + 1] - 1 - count[s[i] + 1]] = i
        count[s[i] + 1] += 1

    count = [0 for _ in range(upper + 1)]
    for i in range(n):
        if sa[i] > 0 and not smaller[sa[i] - 1]:
            sa[bucket[s[sa[i] - 1]] + count[s[sa[i] - 1] + 1]] = sa[i] - 1
            count[s[sa[i] - 1] + 1] += 1

    count = [0 for _ in range(upper + 1)]
    for i in range(n-1, -1, -1):
        if sa[i] > 0 and smaller[sa[i] - 1]:
            sa[bucket[s[sa[i] - 1] + 1] - 1 - count[s[sa[i] - 1] + 1]] = sa[i] - 1
            count[s[sa[i] - 1] + 1] += 1

    new_num = 0
    new_array = [-1 for _ in range(len(lmss))]
    lms_to_ind = dict()
    for i, x in enumerate(lmss):
        lms_to_ind[x] = len(lmss) - 1 - i
    last = sa[0]
    new_array[lms_to_ind[sa[0]]] = 0
    for i in range(1, n):
        if is_lms[sa[i]]:
            tmp, j = sa[i], 0
            while True:
                if s[last + j] != s[tmp + j]:
                    new_num += 1
                    break
                if j > 0 and (is_lms[last + j] or is_lms[tmp + j]):
                    break
                j += 1
            new_array[lms_to_ind[tmp]] = new_num
            last = tmp
            
    if new_num == len(lmss) - 1:
        seed = [-1 for _ in range(len(lmss))]
        for i, x in enumerate(new_array):
            seed[x] = i

    else:
        seed = SA_IS(new_array, upper=new_num + 1)

    new_lmss = [-1 for _ in range(len(lmss))]
    for i, x in enumerate(seed):
        new_lmss[len(lmss) - 1 - i] = lmss[len(lmss) - 1 - x]

    sa = [-1 for _ in range(n)]
    count = [0 for _ in range(upper + 1)]
    for i in new_lmss:
        sa[bucket[s[i] + 1] - 1 - count[s[i] + 1]] = i
        count[s[i] + 1] += 1

    count = [0 for _ in range(upper + 1)]
    for i in range(n):
        if sa[i] > 0 and not smaller[sa[i] - 1]:
            sa[bucket[s[sa[i] - 1]] + count[s[sa[i] - 1] + 1]] = sa[i] - 1
            count[s[sa[i] - 1] + 1] += 1

    count = [0 for _ in range(upper + 1)]
    for i in range(n-1, -1, -1):
        if sa[i] > 0 and smaller[sa[i] - 1]:
            sa[bucket[s[sa[i] - 1] + 1] - 1 - count[s[sa[i] - 1] + 1]] = sa[i] - 1
            count[s[sa[i] - 1] + 1] += 1

    return sa

def suffix_array(s):
    lis = [ord(x)-96 for x in s] + [0]
    sa = SA_IS(lis, 30)[1:]
    return sa

def lcp_array(s, sa):
    n = len(s)
    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 main():
    s = input()[:-1]
    n = len(s)
    sa = suffix_array(s)

    lcp = lcp_array(s, sa)
    
    ans = n * (n+1) // 2 - sum(lcp)
    print(ans)

    
if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip()
sys.setrecursionlimit(max(1000, 10**9))
write = lambda x: sys.stdout.write(x+"\n")


def sa_naive(s):
    n = len(s)
    sa = list(range(n))
    sa.sort(key=lambda x: s[x:])
    return sa

def sa_doubling(s):
    n = len(s)
    sa = list(range(n))
    rnk = s
    tmp = [0] * n
    k = 1
    while k < n:
        sa.sort(key=lambda x: (rnk[x], rnk[x + k])
                if x + k < n else (rnk[x], -1))
        tmp[sa[0]] = 0
        for i in range(1, n):
            tmp[sa[i]] = tmp[sa[i - 1]]
            if sa[i - 1] + k < n:
                x = (rnk[sa[i - 1]], rnk[sa[i - 1] + k])
            else:
                x = (rnk[sa[i - 1]], -1)
            if sa[i] + k < n:
                y = (rnk[sa[i]], rnk[sa[i] + k])
            else:
                y = (rnk[sa[i]], -1)
            if x < y:
                tmp[sa[i]] += 1
        k *= 2
        tmp, rnk = rnk, tmp
    return sa

def sa_is(s, upper):
    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 < 10:
        return sa_naive(s)
    if n < 50:
        return sa_doubling(s)
    ls = [0] * 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 ls[i]:
            sum_l[s[i] + 1] += 1
        else:
            sum_s[s[i]] += 1
    for i in range(upper):
        sum_s[i] += sum_l[i]
        if i < upper:
            sum_l[i + 1] += sum_s[i]
    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 = []
    for i in range(1, n):
        if not ls[i - 1] and ls[i]:
            lms.append(i)
    sa = [-1] * n
    buf = sum_s.copy()
    for d in lms:
        if d == n:
            continue
        sa[buf[s[d]]] = d
        buf[s[d]] += 1
    buf = sum_l.copy()
    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 = sum_l.copy()
    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
    if m:
        sorted_lms = []
        for v in sa:
            if lms_map[v] != -1:
                sorted_lms.append(v)
        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]]
        sa = [-1] * n
        buf = sum_s.copy()
        for d in sorted_lms:
            if d == n:
                continue
            sa[buf[s[d]]] = d
            buf[s[d]] += 1
        buf = sum_l.copy()
        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 = sum_l.copy()
        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
    return sa

def suffix_array(s, upper=255):
    if type(s) is str:
        s = [ord(c) for c in s]
    return sa_is(s, upper)

def lcp_array(s, sa):
    n = len(s)
    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

s = input()
n = len(s)
sa = suffix_array(s)
lcp = lcp_array(s,sa)
ans = n*(n+1)//2 - sum(lcp)
print(ans)

Submission status: Accepted

Solution code:
import sys
import numpy as np
import numba
from numba import njit, b1, i4, i8, f8

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

@njit((i8[:], ), cache=True)
def suffix_array(S):
    """ sa by doubling. O(N log^2N)"""
    N = len(S)
    SA = np.arange(N, dtype=np.int64)
    rnk = S
    tmp = np.empty(N, np.int64)
    k = 1
    while True:
        comp_key = np.empty(N, np.int64)
        for i in range(N):
            if i + k < N:
                comp_key[i] = rnk[i] << 32 | (1 + rnk[i + k])
            else:
                comp_key[i] = rnk[i] << 32
        SA = comp_key.argsort()
        rnk = np.empty(N, np.int64)
        rnk[SA[0]] = 0
        for i in range(1, N):
            rnk[SA[i]] = rnk[SA[i -
                                1]] + (comp_key[SA[i - 1]] < comp_key[SA[i]])
        k *= 2
        if k >= N:
            break
    return SA


@njit((i8[:], i8[:]), cache=True)
def lcp_array(S, SA):
    N = len(S)
    rnk = np.empty(N, np.int64)
    for i in range(N):
        rnk[SA[i]] = i
    LCP = np.empty(N - 1, np.int64)
    h = 0
    for i in range(N):
        if h:
            h -= 1
        if not rnk[i]:
            continue
        j = SA[rnk[i] - 1]
        while True:
            if S[j + h] != S[i + h]:
                break
            h += 1
            if i + h >= N or j + h >= N:
                break
        LCP[rnk[i] - 1] = h
    return LCP

def main(S):
    N = len(S)
    SA = suffix_array(S)
    LCP = lcp_array(S, SA)
    return N * (N + 1) // 2 - LCP.sum()

S = np.array(list(read().rstrip()), np.int64)

print(main(S))

Submission status: Accepted

Solution code:


class StrAlg:
    @staticmethod
    def sa_naive(s):
        n = len(s)
        sa = list(range(n))
        sa.sort(key=lambda x: s[x:])
        return sa

    @staticmethod
    def sa_doubling(s):
        n = len(s)
        sa = list(range(n))
        rnk = s
        tmp = [0] * n
        k = 1
        while k < n:
            sa.sort(key=lambda x: (rnk[x], rnk[x + k])
                    if x + k < n else (rnk[x], -1))
            tmp[sa[0]] = 0
            for i in range(1, n):
                tmp[sa[i]] = tmp[sa[i - 1]]
                if sa[i - 1] + k < n:
                    x = (rnk[sa[i - 1]], rnk[sa[i - 1] + k])
                else:
                    x = (rnk[sa[i - 1]], -1)
                if sa[i] + k < n:
                    y = (rnk[sa[i]], rnk[sa[i] + k])
                else:
                    y = (rnk[sa[i]], -1)
                if x < y:
                    tmp[sa[i]] += 1
            k *= 2
            tmp, rnk = rnk, tmp
        return sa

    @staticmethod
    def sa_is(s, upper):
        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 < 10:
            return StrAlg.sa_naive(s)
        if n < 50:
            return StrAlg.sa_doubling(s)
        ls = [0] * 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 ls[i]:
                sum_l[s[i] + 1] += 1
            else:
                sum_s[s[i]] += 1
        for i in range(upper):
            sum_s[i] += sum_l[i]
            if i < upper:
                sum_l[i + 1] += sum_s[i]
        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 = []
        for i in range(1, n):
            if not ls[i - 1] and ls[i]:
                lms.append(i)
        sa = [-1] * n
        buf = sum_s.copy()
        for d in lms:
            if d == n:
                continue
            sa[buf[s[d]]] = d
            buf[s[d]] += 1
        buf = sum_l.copy()
        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 = sum_l.copy()
        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
        if m:
            sorted_lms = []
            for v in sa:
                if lms_map[v] != -1:
                    sorted_lms.append(v)
            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 = StrAlg.sa_is(rec_s, rec_upper)
            for i in range(m):
                sorted_lms[i] = lms[rec_sa[i]]
            sa = [-1] * n
            buf = sum_s.copy()
            for d in sorted_lms:
                if d == n:
                    continue
                sa[buf[s[d]]] = d
                buf[s[d]] += 1
            buf = sum_l.copy()
            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 = sum_l.copy()
            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
        return sa

    @staticmethod
    def suffix_array(s, upper=255):
        if type(s) is str:
            s = [ord(c) for c in s]
        return StrAlg.sa_is(s, upper)

    @staticmethod
    def lcp_array(s, sa):
        n = len(s)
        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

    @staticmethod
    def z_algorithm(s):
        n = len(s)
        if n == 0:
            return []
        z = [0] * n
        j = 0
        for i in range(1, n):
            z[i] = 0 if j + z[j] <= i else min(j + z[j] - i, z[i - j])
            while i + z[i] < n and s[z[i]] == s[i + z[i]]:
                z[i] += 1
                if j + z[j] < i + z[i]:
                    j = i
        z[0] = n
        return z


def atcoder_practice2_i():
    S = input()
    n = len(S)
    ans = (n+1)*n // 2
    sa = StrAlg.suffix_array(S)
    ans -= sum(StrAlg.lcp_array(S, sa))
    print(ans)


if __name__ == "__main__":
    atcoder_practice2_i()

Submission status: Accepted

Solution code:


class StrAlg:
    @staticmethod
    def sa_naive(s):
        n = len(s)
        sa = list(range(n))
        sa.sort(key=lambda x: s[x:])
        return sa

    @staticmethod
    def sa_doubling(s):
        n = len(s)
        sa = list(range(n))
        rnk = s
        tmp = [0] * n
        k = 1
        while k < n:
            sa.sort(key=lambda x: (rnk[x], rnk[x + k])
                    if x + k < n else (rnk[x], -1))
            tmp[sa[0]] = 0
            for i in range(1, n):
                tmp[sa[i]] = tmp[sa[i - 1]]
                if sa[i - 1] + k < n:
                    x = (rnk[sa[i - 1]], rnk[sa[i - 1] + k])
                else:
                    x = (rnk[sa[i - 1]], -1)
                if sa[i] + k < n:
                    y = (rnk[sa[i]], rnk[sa[i] + k])
                else:
                    y = (rnk[sa[i]], -1)
                if x < y:
                    tmp[sa[i]] += 1
            k *= 2
            tmp, rnk = rnk, tmp
        return sa

    @staticmethod
    def sa_is(s, upper):
        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 < 10:
            return StrAlg.sa_naive(s)
        if n < 50:
            return StrAlg.sa_doubling(s)
        ls = [0] * 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 ls[i]:
                sum_l[s[i] + 1] += 1
            else:
                sum_s[s[i]] += 1
        for i in range(upper):
            sum_s[i] += sum_l[i]
            if i < upper:
                sum_l[i + 1] += sum_s[i]
        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 = []
        for i in range(1, n):
            if not ls[i - 1] and ls[i]:
                lms.append(i)
        sa = [-1] * n
        buf = sum_s.copy()
        for d in lms:
            if d == n:
                continue
            sa[buf[s[d]]] = d
            buf[s[d]] += 1
        buf = sum_l.copy()
        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 = sum_l.copy()
        for i in range(n)[::-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
        if m:
            sorted_lms = []
            for v in sa:
                if lms_map[v] != -1:
                    sorted_lms.append(v)
            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 = StrAlg.sa_is(rec_s, rec_upper)
            for i in range(m):
                sorted_lms[i] = lms[rec_sa[i]]
            sa = [-1] * n
            buf = sum_s.copy()
            for d in sorted_lms:
                if d == n:
                    continue
                sa[buf[s[d]]] = d
                buf[s[d]] += 1
            buf = sum_l.copy()
            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 = sum_l.copy()
            for i in range(n)[::-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
        return sa

    @staticmethod
    def suffix_array(s, upper=255):
        if type(s) is str:
            s = [ord(c) for c in s]
        return StrAlg.sa_is(s, upper)

    @staticmethod
    def lcp_array(s, sa):
        n = len(s)
        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

    @staticmethod
    def z_algorithm(s):
        n = len(s)
        if n == 0:
            return []
        z = [0] * n
        j = 0
        for i in range(1, n):
            z[i] = 0 if j + z[j] <= i else min(j + z[j] - i, z[i - j])
            while i + z[i] < n and s[z[i]] == s[i + z[i]]:
                z[i] += 1
                if j + z[j] < i + z[i]:
                    j = i
        z[0] = n
        return z


def atcoder_practice2_i():
    S = input()
    n = len(S)
    ans = (n+1)*n // 2
    sa = StrAlg.suffix_array(S)
    ans -= sum(StrAlg.lcp_array(S, sa))
    print(ans)


if __name__ == "__main__":
    atcoder_practice2_i()