Delicious Dessert (Codeforces 1780G)
Tags: Math, string suffix structures

Today is an important day for chef Tonio — an auditor has arrived in his hometown of Morioh. He has also arrived at Tonio's restaurant and ordered dessert. Tonio has not been prepared for this turn of events.As you know, dessert is a string of lowercase English letters. Tonio remembered the rule of desserts — a string $$$s$$$ of length $$$n$$$. Any dessert $$$t$$$ is delicious if the number of occurrences of $$$t$$$ in $$$s$$$ as a substring is divisible by the length of $$$t$$$.Now Tonio wants to know the number of delicious substrings of $$$s$$$. If the substring occurs several times in the string $$$s$$$, then all occurrences must be taken into account.

Input: The first line contains an integer $$$n$$$ ($$$1 \leq n \leq 10^6$$$) — the length of the rule $$$s$$$.The second line contains the string $$$s$$$ of length $$$n$$$ — the rule. The rule consists only of lowercase English letters.

Output: In a single line print the number of delicious substrings of $$$s$$$.

Note: In the first sample, there are many delicious substrings. $$$7$$$ of them are substrings of length $$$1$$$ (because any number is divisible by $$$1$$$). Consider other delicious substrings:  "ab" occurs in $$$s$$$ $$$2$$$ times, which is divisible by the length of the substring.  "ba" also occurs $$$2$$$ times. Therefore, the answer is $$$7 + 2 + 2 = 11$$$.Note that the answer includes both occurrences of "ab" and "ba".

Sample Inputs:
7
abacaba

8
abaababa

10
deadinside

5
aaaaa

Sample Outputs:
11

11

12

12

Solution (Python):
import sys, random
input = lambda : sys.stdin.readline().rstrip()
write = lambda x: sys.stdout.write(x + '\n')
writef = lambda x: print('{:.12f}'.format(x))
debug = lambda x: sys.stderr.write(x + '\n')
YES = 'Yes'
NO = 'No'
pans = lambda v: print(YES if v else NO)
INF = 10 ** 18
LI = lambda : list(map(int, input().split()))
II = lambda : int(input())
SI = lambda : [ord(c) - ord('a') for c in input()]

def debug(_l_):
    for s in _l_.split():
        print(f'{s}={eval(s)}', end=' ')
    print()

def dlist(*l, fill=0):
    if len(l) == 1:
        return [fill] * l[0]
    ll = l[1:]
    return [dlist(*ll, fill=fill) for _ in range(l[0])]

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):
    if type(s) is str:
        s = [ord(c) for c in s]
    assert min(s) >= 0
    return sa_is(s[:], max(s))

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

class SG:

    def __init__(self, n, v=None):
        self._n = n
        self.geta = 0
        x = 0
        while 1 << x < n:
            x += 1
        self._log = x
        self._size = 1 << self._log
        self._d = [ninf] * (2 * self._size)
        if v is not None:
            for i in range(self._n):
                self._d[self._size + i] = v[i]
        for i in range(self._size - 1, 0, -1):
            self._update(i)

    def _update(self, k):
        self._d[k] = op(self._d[2 * k], self._d[2 * k + 1])

    def update(self, p, x):
        assert 0 <= p < self._n
        p += self._size
        self._d[p] = x
        for i in range(1, self._log + 1):
            k = p >> i
            self._d[k] = op(self._d[2 * k], self._d[2 * k + 1])

    def get(self, p):
        assert 0 <= p < self._n
        return self._d[p + self._size]

    def check(self):
        return [self.get(p) for p in range(self._n)]

    def query(self, left, right):
        assert 0 <= left <= right <= self._n
        sml = ninf
        smr = ninf
        left += self._size
        right += self._size
        while left < right:
            if left & 1:
                sml = op(sml, self._d[left])
                left += 1
            if right & 1:
                right -= 1
                smr = op(self._d[right], smr)
            left >>= 1
            right >>= 1
        return op(sml, smr)

    def query_all(self):
        return self._d[1]

    def max_right(self, left, f):
        """"""
        if left == self._n:
            return self._n
        left += self._size
        sm = ninf
        first = True
        while first or left & -left != left:
            first = False
            while left % 2 == 0:
                left >>= 1
            if not f(op(sm, self._d[left])):
                while left < self._size:
                    left *= 2
                    if f(op(sm, self._d[left])):
                        sm = op(sm, self._d[left])
                        left += 1
                return left - self._size
            sm = op(sm, self._d[left])
            left += 1
        return self._n

    def min_left(self, right, f):
        """"""
        if right == 0:
            return 0
        right += self._size
        sm = ninf
        first = True
        while first or right & -right != right:
            first = False
            right -= 1
            while right > 1 and right % 2:
                right >>= 1
            if not f(op(self._d[right], sm)):
                while right < self._size:
                    right = 2 * right + 1
                    if f(op(self._d[right], sm)):
                        sm = op(self._d[right], sm)
                        right -= 1
                return right + 1 - self._size
            sm = op(self._d[right], sm)
        return 0
op = min
ninf = 10 ** 12
n = int(input())
s = SI()
sa = suffix_array(s)
rank = [0] * len(sa)
for (i, v) in enumerate(sa):
    rank[v] = i
lcp = lcp_array(s, sa)
ans = 0
sg = SG(len(lcp), lcp)
p = 0
for i in range(n):
    if i > 0:
        p = lcp[i - 1]
    ind = sa[i]
    for v in range(p + 1, n - ind + 1):
        res = sg.max_right(i, lambda val: val >= v)
        if res - i + 1 < v:
            break
        if (res - i + 1) % v == 0:
            ans += res - i + 1
print(ans)

Solution (Python):
import sys, random
input = lambda : sys.stdin.readline().rstrip()
write = lambda x: sys.stdout.write(x + '\n')
writef = lambda x: print('{:.12f}'.format(x))
debug = lambda x: sys.stderr.write(x + '\n')
YES = 'Yes'
NO = 'No'
pans = lambda v: print(YES if v else NO)
INF = 10 ** 18
LI = lambda : list(map(int, input().split()))
II = lambda : int(input())
SI = lambda : [ord(c) - ord('a') for c in input()]

def debug(_l_):
    for s in _l_.split():
        print(f'{s}={eval(s)}', end=' ')
    print()

def dlist(*l, fill=0):
    if len(l) == 1:
        return [fill] * l[0]
    ll = l[1:]
    return [dlist(*ll, fill=fill) for _ in range(l[0])]

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):
    if type(s) is str:
        s = [ord(c) for c in s]
    assert min(s) >= 0
    return sa_is(s[:], max(s))

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

class SG:

    def __init__(self, n, v=None):
        self._n = n
        self.geta = 0
        x = 0
        while 1 << x < n:
            x += 1
        self._log = x
        self._size = 1 << self._log
        self._d = [ninf] * (2 * self._size)
        if v is not None:
            for i in range(self._n):
                self._d[self._size + i] = v[i]
        for i in range(self._size - 1, 0, -1):
            self._update(i)

    def _update(self, k):
        self._d[k] = op(self._d[2 * k], self._d[2 * k + 1])

    def update(self, p, x):
        assert 0 <= p < self._n
        p += self._size
        self._d[p] = x
        for i in range(1, self._log + 1):
            k = p >> i
            self._d[k] = op(self._d[2 * k], self._d[2 * k + 1])

    def get(self, p):
        assert 0 <= p < self._n
        return self._d[p + self._size]

    def check(self):
        return [self.get(p) for p in range(self._n)]

    def query(self, left, right):
        assert 0 <= left <= right <= self._n
        sml = ninf
        smr = ninf
        left += self._size
        right += self._size
        while left < right:
            if left & 1:
                sml = op(sml, self._d[left])
                left += 1
            if right & 1:
                right -= 1
                smr = op(self._d[right], smr)
            left >>= 1
            right >>= 1
        return op(sml, smr)

    def query_all(self):
        return self._d[1]

    def max_right(self, left, f):
        """"""
        if left == self._n:
            return self._n
        left += self._size
        sm = ninf
        first = True
        while first or left & -left != left:
            first = False
            while left % 2 == 0:
                left >>= 1
            if not f(op(sm, self._d[left])):
                while left < self._size:
                    left *= 2
                    if f(op(sm, self._d[left])):
                        sm = op(sm, self._d[left])
                        left += 1
                return left - self._size
            sm = op(sm, self._d[left])
            left += 1
        return self._n

    def min_left(self, right, f):
        """"""
        if right == 0:
            return 0
        right += self._size
        sm = ninf
        first = True
        while first or right & -right != right:
            first = False
            right -= 1
            while right > 1 and right % 2:
                right >>= 1
            if not f(op(self._d[right], sm)):
                while right < self._size:
                    right = 2 * right + 1
                    if f(op(self._d[right], sm)):
                        sm = op(self._d[right], sm)
                        right -= 1
                return right + 1 - self._size
            sm = op(self._d[right], sm)
        return 0
op = min
ninf = 10 ** 12
n = int(input())
s = SI()
sa = suffix_array(s)
rank = [0] * len(sa)
for (i, v) in enumerate(sa):
    rank[v] = i
lcp = lcp_array(s, sa)
ans = 0
sg = SG(len(lcp), lcp)
p = 0
for i in range(n):
    if i > 0:
        p = lcp[i - 1]
    ind = sa[i]
    for v in range(p + 1, n - ind + 1):
        res = sg.max_right(i, lambda val: val >= v)
        if res - i + 1 < v:
            break
        if (res - i + 1) % v == 0:
            ans += res - i + 1
print(ans)

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def suffix_array(s):
    b = (len(s) - 1).bit_length()
    u = list(s)
    (l, j) = (len(u), 1)
    (v, w) = ([0] * l, [0] * l)
    for _ in range(b):
        m = max(u)
        s0 = [0] * (m + 3)
        for i in u:
            s0[i + 2] += 1
        for i in range(1, m + 3):
            s0[i] += s0[i - 1]
        for i in range(l):
            k = u[i] + 1
            v[s0[k]] = i - j
            s0[k] += 1
        for i in v:
            k = u[i]
            w[s0[k]] = i
            s0[k] += 1
        (c, la1, la2) = (0, 0, 0)
        for i in w:
            (x, y) = (u[i], u[(i + j) % l])
            if la1 ^ x or la2 ^ y:
                (la1, la2) = (x, y)
                c += 1
            v[i] = c
        (u, v) = (v, u)
        if c == l:
            break
        j <<= 1
    l = u[-1]
    sa = [0] * len(s)
    for i in range(len(s) - 1):
        sa[u[i] - l] = i
    sa[0] = len(sa) - 1
    return sa

def lcp_array(s, sa):
    l = len(s)
    (lcp, rank) = ([0] * l, [0] * l)
    for i in range(l):
        rank[sa[i]] = i
    v = 0
    for i in range(l):
        if not rank[i]:
            continue
        j = sa[rank[i] - 1]
        u = min(i, j)
        while u + v < l and s[i + v] == s[j + v]:
            v += 1
        lcp[rank[i]] = v
        v = max(v - 1, 0)
    return lcp

def make_graph(n, m):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, 2):
        s[x[i] + 2] += 1
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * m
    for i in range(0, 2 * m, 2):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def get_root(s):
    if s ^ root[s]:
        root[s] = get_root(root[s])
        return root[s]
    return s

def unite(s, t):
    (s, t) = (get_root(s), get_root(t))
    if s == t:
        return
    if rank[s] < rank[t]:
        (s, t) = (t, s)
    if rank[s] == rank[t]:
        rank[s] += 1
    root[t] = s
    size[s] += size[t]
    return

def same(s, t):
    return True if get_root(s) == get_root(t) else False

def get_size(s):
    return size[get_root(s)]
n = int(input())
s = list(input().rstrip())
s = [i - 96 for i in s] + [0]
sa = suffix_array(s)
lcp = lcp_array(s, sa)
x = [0] * (2 * (n - 1))
for i in range(n - 1):
    (x[2 * i], x[2 * i + 1]) = (lcp[i + 2], i + 1)
(G, s0) = make_graph(n, n - 1)
cnt = [0] * (n + 1)
cnt[1] = n
root = [i for i in range(n + 1)]
rank = [1 for _ in range(n + 1)]
size = [1 for _ in range(n + 1)]
ans = 0
for i in range(n, 0, -1):
    for v in range(s0[i], s0[i + 1]):
        j = G[v]
        cnt[get_size(j)] -= 1
        cnt[get_size(j + 1)] -= 1
        unite(j, j + 1)
        cnt[get_size(j)] += 1
    for j in range(i, n + 1, i):
        ans += cnt[j] * j
print(ans)