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: Wrong Answer

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


def main():
    import sys
    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):
            P_new = []
            C_new = [0] * N
            cnt = [[] for _ in range(N)]
            for i in range(N):
                P[i] = (P[i] - (1 << k)) % N
                cnt[C[P[i]]].append(P[i])
            for i_list in cnt:
                for i in i_list:
                    P_new.append(i)
            cnt_c = -1
            prev = -1
            for p in P_new:
                new = C[p] * N + C[(p + (1 << k)) % N]
                if new != prev:
                    cnt_c += 1
                    prev = new
                C_new[p] = cnt_c
            P = P_new
            C = C_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: Wrong Answer

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 - 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(StrAlg.z_algorithm(S))
    print(ans)


if __name__ == "__main__":
    atcoder_practice2_i()

Submission status: Wrong Answer

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]]
                x = (rnk[sa[i - 1]], rnk[sa[i - 1] + k]) if sa[i - 1] + \
                    k < n else (rnk[sa[i - 1]], -1)
                y = (rnk[sa[i]], rnk[sa[i] + k]) if sa[i] + \
                    k < n else (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 - 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)
        #assert n >= 1
        s2 = [ord(c) for c in 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)
        s2 = [ord(c) for c in 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))

if __name__ == "__main__":
    atcoder_practice2_i()

Submission status: Wrong Answer

Solution code:
import sys
input = sys.stdin.readline
S = input()[: -1]
N = len(S)

def SuffixArray(n, s):
  s += "??"
  s0 = []
  s1 = []
  s2 = []
  for i in range(n):
    if i % 3 == 1: s1.append((s[i: i + 3], i))
    if i % 3 == 2: s2.append((s[i: i + 3], i))
  s12 = s1 + s2
  s12.sort(key = lambda x: (x[0], -x[1]))
  table = [-1] * (n + 2)
  for i in range(len(s12)): table[s12[i][1]] = i
  for i in range(n):
    if i % 3 == 0: s0.append((s[i], table[i + 1], i))
  s0.sort(key = lambda x: (x[0], x[1]))
  sa = []
  s0.reverse()
  s12.reverse()
  while len(sa) < n:
    if len(s0) == 0:
      sa.append(s12.pop()[1])
      continue
    if len(s12) == 0:
      sa.append(s0.pop()[2])
      continue
    c, r, i = s0.pop()
    x, j = s12.pop()
    cc = s[i + 1]
    rr = table[i + 2]
    xx = s[j + 1]
    xr = table[j + 1]
    xrr = table[j + 2]
    if j % 3 == 1:
      if c < x[0]:
        sa.append(i)
        s12.append((x, j))
      if c > x[0]:
        sa.append(j)
        s0.append((c, r, i))
      if c == x[0]:
        if r < xr:
          sa.append(i)
          s12.append((x, j))
        if r > xr:
          sa.append(j)
          s0.append((c, r, i))
    if j % 3 == 2:
      if c < x[0]:
        sa.append(i)
        s12.append((x, j))
      if c > x[0]:
        sa.append(j)
        s0.append((c, r, i))
      if c == x[0]:
        if cc < xx:
          sa.append(i)
          s12.append((x, j))
        if cc > xx:
          sa.append(j)
          s0.append((c, r, i))
        if cc == xx:
          if rr < xrr:
            sa.append(i)
            s12.append((x, j))
          if rr > xrr:
            sa.append(j)
            s0.append((c, r, i))
  return sa

def LCPArray(n, s, sa):
  r = [0] * n
  for i in range(n): r[sa[i]] = i
  lcp = [0] * (n - 1)
  x = 0
  for i in range(n):
    if x > 0: x -= 1
    if r[i] == 0: continue
    j = sa[r[i] - 1]
    while max(i, j) + x < n:
      if s[i + x] != s[j + x]: break
      x += 1
    lcp[r[i] - 1] = x
  return lcp 
sa = SuffixArray(N, S)
lcp = LCPArray(N, S, sa)
print(N * (N + 1) // 2 - sum(lcp))