Problem set

Why this matters when you build with AI

String code is easy to get almost right. A generated program often passes the sample, which is short, and then fails on a case the sample never had: a pattern as long as the text, a string of one repeated letter, a tie between two rotations, or an empty range. Samples are too small to show these.

This set is a place to practise the habit that catches them. For each problem you pick a tool from the module, write a brute force that is too slow for large input but obviously right, and compare the two on many small random strings with a small alphabet. A small alphabet matters: with only a and b, repeats and ties appear often, so the random tests reach the hard cases quickly.

How to work the set

Give yourself time for each problem and move on if you are stuck after a while. For each problem, do these steps in order.

1. Read the statement and write the limits on paper. Note the largest value that can appear.
2. Pick the tool. Write one sentence on why it fits, and one on what it cannot do.
3. Write a brute force first. Keep it simple enough to trust.
4. Write the fast solution. Use 64-bit integers for counts and sums.
5. Run both on the sample, then on at least 500 random strings of length 1 to 14 over the letters a, b and c.
6. Time your fast solution at the largest limit. Compare the time with the target in the statement.

The answers below give the method, a Python reference, and what we checked. The Python references were run on small random inputs against a brute force. We did not write C++ versions for this set, and we did not time anything at the largest limits. Aim for the limits in the statements when you write your own C++.

Helpers used in several answers

These three functions appear in more than one answer. The prefix function is covered in the prefix function lesson. The hash is a polynomial hash modulo 2^61 - 1 with a fixed base. A fixed base is acceptable for this set, but a contest solution should draw a random base, because a fixed base can be attacked with a crafted input. The Manacher function returns the radius of the palindromes around each centre of the string written with # between letters.

def pf(s):
    pi = [0] * len(s)
    for i in range(1, len(s)):
        k = pi[i - 1]
        while k > 0 and s[i] != s[k]:
            k = pi[k - 1]
        if s[i] == s[k]:
            k += 1
        pi[i] = k
    return pi
 
 
MOD = (1 << 61) - 1
BASE = 1000003
 
 
def hash_table(s):
    h = [0] * (len(s) + 1)
    pw = [1] * (len(s) + 1)
    for i, ch in enumerate(s):
        h[i + 1] = (h[i] * BASE + ord(ch)) % MOD
        pw[i + 1] = pw[i] * BASE % MOD
    return h, pw
 
 
def sub_hash(h, pw, l, r):
    # hash of s[l..r) with the polynomial hash above
    return (h[r] - h[l] * pw[r - l]) % MOD
 
 
def manacher_radii(s):
    # p[i] = radius of the palindrome centred at i in t = '#' + '#'.join(s) + '#'
    t = "#" + "#".join(s) + "#"
    n = len(t)
    p = [0] * n
    c = r = 0
    for i in range(n):
        if i < r:
            p[i] = min(r - i, p[2 * c - i])
        while i - p[i] - 1 >= 0 and i + p[i] + 1 < n and t[i - p[i] - 1] == t[i + p[i] + 1]:
            p[i] += 1
        if i + p[i] > r:
            c, r = i, i + p[i]
    return p

The problems

Problem 1: smallest period of every prefix

Statement. Given a string s, for each prefix of s print its smallest period. A prefix of length L has smallest period p when p is the smallest positive number with s[i] = s[i + p] for every i where both indexes lie inside the prefix.

Input. One line with s.

Output. n numbers on one line, separated by spaces. The L-th number belongs to the prefix of length L.

Limits. 1 ≤ n ≤ 200000. Lowercase letters only.

Example.

input
aabaab
output
1 1 3 3 3 3

Hint. Look at what the last value of the prefix function says about a prefix, and what a border has to do with a period.

Target. O(n).

Show the answer

Use the prefix function. The smallest period of the prefix of length L is L - pi[L - 1]. The Python reference below is checked on 3000 random strings of length 1 to 14 over a and b, against a brute force that tries every shift. The sample gives 1 1 3 3 3 3.

def p1_fast(s):
    pi = pf(s)
    return [L - pi[L - 1] for L in range(1, len(s) + 1)]

Problem 2: prefixes that are powers

Statement. Count the prefixes of s (lengths 1 to n) that are a power of a block with exponent at least 2. A string is such a power when it equals u written k times in a row for some non-empty u and some k at least 2.

Input. One line with s.

Output. One number.

Limits. 1 ≤ n ≤ 200000. Lowercase letters only.

Example.

input
aaaa
output
3

The prefixes aa, aaa and aaaa are powers. The prefix a is not, because k must be at least 2. The second example is abab, which gives 1, because only the whole string is a power.

Hint. The question depends only on the smallest period of each prefix and on which lengths that period divides. The periodicity lesson has the rule.

Target. O(n).

Show the answer

For each prefix of length L, let p = L - pi[L - 1]. The prefix is a power with exponent at least 2 exactly when p < L and p divides L. The reference is checked on 3000 random strings of length 1 to 14 over a and b, against a brute force that tries every proper block length. Sample results: aaaa gives 3, abab gives 1.

def p2_fast(s):
    pi = pf(s)
    count = 0
    for L in range(1, len(s) + 1):
        p = L - pi[L - 1]
        if p < L and L % p == 0:
            count += 1
    return count

Problem 3: occurrences of a pattern, overlaps counted

Statement. Given a pattern p and a text t, count the positions i where p occurs in t starting at i. Occurrences may overlap.

Input. Two lines: p, then t.

Output. One number.

Limits. 1 ≤ |p| ≤ |t| ≤ 1000000. Lowercase letters only.

Example.

input
aba
abababa
output
3

The occurrences start at positions 0, 2 and 4.

Hint. One pass over the text is enough if you keep the longest match you have so far. A character that never appears in either string can separate them in one combined string.

Target. O(|p| + |t|).

Show the answer

Build the combined string p + "#" + t, where # is not a letter, and run the prefix function on it. Each position where the value equals |p| is one occurrence. The # stops any border from crossing the separator. The reference is checked on 3000 random pairs with p of length 1 to 4 and t of length 4 to 14 over a and b, against a brute force that compares every window. Sample: aba in abababa gives 3.

def p3_fast(p, t):
    pi = pf(p + "#" + t)
    m = len(p)
    return sum(1 for v in pi[m + 1:] if v == m)

Problem 4: the shortest string that contains two strings

Statement. Given a and b, print the length of the shortest string that contains both a and b as substrings.

Input. Two lines: a, then b.

Output. One number.

Limits. 1 ≤ |a|, |b| ≤ 100000. Lowercase letters only.

Example.

input
aba
bab
output
4

The string abab contains aba at position 0 and bab at position 1. The second example is abc and cde, which give 5 (abcde).

Hint. The shortest answer overlaps the two strings as much as it can. Ask which suffix of one string can also be a prefix of the other.

Target. O(|a| + |b|).

Show the answer

If one string already contains the other, the answer is the length of the longer one. Otherwise, let k1 be the length of the longest prefix of b that is also a suffix of a, and k2 the same with a and b swapped. The answer is the smaller of |a| + |b| - k1 and |a| + |b| - k2. The prefix function on b + "#" + a gives k1 in one pass. The reference is checked on 400 random pairs of length 1 to 4 over a, b and c, against a brute force that tries every string over those letters, in order of length. Samples: aba, bab give 4, and abc, cde give 5.

def p4_fast(a, b):
    if b in a:
        return len(a)
    if a in b:
        return len(b)
    k1 = pf(b + "#" + a)[-1]  # longest prefix of b that is a suffix of a
    k2 = pf(a + "#" + b)[-1]  # longest prefix of a that is a suffix of b
    return min(len(a) + len(b) - k1, len(a) + len(b) - k2)

Problem 5: distinct substrings of one length

Statement. Given s and a number L, count the different substrings of s that have length exactly L. Two substrings are the same when they have the same letters.

Input. Line 1: s. Line 2: L.

Output. One number.

Limits. 1 ≤ L ≤ n ≤ 200000.

Example.

input
ababa
2
output
2

The substrings of length 2 are ab, ba, ab and ba. There are two different ones.

Hint. Equal substrings must get equal keys, and comparing keys must be cheap. A number computed from the letters can serve as a key.

Target. O(n) expected.

Show the answer

Give every window of length L a hash, using the helper above, and count the distinct hash values. A hash can collide, so a contest solution should use two moduli or a 64-bit modulus with a random base. The reference is checked on 2000 random strings of length 1 to 14 over a and b, with a random L, against a brute force that builds the set of substrings. Sample: ababa, L = 2 gives 2.

def p5_fast(s, L):
    if L > len(s):
        return 0
    h, pw = hash_table(s)
    return len({sub_hash(h, pw, i, i + L) for i in range(len(s) - L + 1)})

Problem 6: the longest common substring

Statement. Given a and b, print the length of the longest string that is a substring of both.

Input. Two lines: a, then b.

Output. One number.

Limits. 1 ≤ |a|, |b| ≤ 2000. Lowercase letters only.

Example.

input
xabcdy
zabcq
output
3

The common substring abc has length 3.

Hint. If a common substring of length m exists, do shorter common substrings exist too? If yes, you can search for the length.

Target. O(n log n) with hashing, or O(|a| · |b|) with a dynamic program.

Show the answer

If a common substring of length m exists, a common substring of every length below m exists too, because any part of it is also common. So the answer can be found by binary search on the length. For a length m, put the hashes of all windows of a in a set, and look for a window of b in that set. The reference is checked on 2000 random pairs of length 1 to 12 over a and b, against a brute force that extends each pair of starts. Sample: xabcdy, zabcq gives 3.

def p6_fast(a, b):
    ha, pa = hash_table(a)
    hb, pb = hash_table(b)
 
    def ok(m):
        if m == 0:
            return True
        seen = {sub_hash(ha, pa, i, i + m) for i in range(len(a) - m + 1)}
        return any(sub_hash(hb, pb, j, j + m) in seen for j in range(len(b) - m + 1))
 
    lo, hi = 0, min(len(a), len(b))
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if ok(mid):
            lo = mid
        else:
            hi = mid - 1
    return lo

Problem 7: is each query a substring

Statement. Given s, then q strings x_1 to x_q, print YES if x_i is a substring of s, and NO otherwise.

Input. Line 1: s. Line 2: q. Then q lines, one string each.

Output. q lines, each YES or NO.

Limits. 1 ≤ |s| ≤ 100000, 1 ≤ q ≤ 100000, and the total length of the queries is at most 100000. Lowercase letters only.

Example.

input
banana
3
ana
nab
b
output
YES
NO
YES

Hint. You need one structure built from s that accepts a query by walking through it letter by letter, and that has at most a linear number of states.

Target. O(|s| + total query length).

Show the answer

Build a suffix automaton on s, the structure from the suffix automaton lesson. Walk each query from the start state, following one letter at a time. If a letter has no edge, the answer is NO. If the walk ends, the answer is YES. The reference is checked on 1000 random cases with s of length 1 to 14 and 6 queries of length 1 to 5 over a and b, against a brute force that uses the in test. Sample: the example above gives YES, NO, YES.

def build_sam(s):
    link, length, nxt = [-1], [0], [{}]
    last = 0
    for ch in s:
        cur = len(length)
        length.append(length[last] + 1)
        link.append(-1)
        nxt.append({})
        p = last
        while p != -1 and ch not in nxt[p]:
            nxt[p][ch] = cur
            p = link[p]
        if p == -1:
            link[cur] = 0
        else:
            q = nxt[p][ch]
            if length[p] + 1 == length[q]:
                link[cur] = q
            else:
                clone = len(length)
                length.append(length[p] + 1)
                link.append(link[q])
                nxt.append(dict(nxt[q]))
                while p != -1 and nxt[p].get(ch) == q:
                    nxt[p][ch] = clone
                    p = link[p]
                link[q] = clone
                link[cur] = clone
        last = cur
    return nxt
 
 
def p7_fast(s, queries):
    nxt = build_sam(s)
    out = []
    for x in queries:
        state = 0
        ok = True
        for ch in x:
            if ch not in nxt[state]:
                ok = False
                break
            state = nxt[state][ch]
        out.append("YES" if ok else "NO")
    return out

Problem 8: palindromes of length at least 3

Statement. Count the palindromic substrings of s whose length is at least 3. A palindrome reads the same forwards and backwards.

Input. One line with s.

Output. One number.

Limits. 1 ≤ n ≤ 100000. Lowercase letters only.

Example.

input
aabaa
output
2

The two palindromes are aba and aabaa. The substring aa has length 2, so it is not counted.

Hint. Every palindrome has a centre. Count the palindromes for each centre, then remove the ones you must not count. Treat odd and even centres separately, or use the transformed string.

Target. O(n).

Show the answer

Run Manacher's algorithm on the string written with # between letters. For each centre, the radius tells how many palindromes have that centre. Odd palindromes of length 3 and up are counted by the radius at the letter centres. Even palindromes are counted at the # centres, and length 2 must be removed from the even count. The reference is checked on 2000 random strings of length 1 to 16 over a and b, against a brute force that tests every substring. Sample: aabaa gives 2.

def p8_fast(s):
    p = manacher_radii(s)
    odd = sum(p[i] // 2 for i in range(1, len(p), 2))  # odd palindromes of length 3 and up
    even = sum((p[i] + 1) // 2 for i in range(2, len(p), 2))  # even palindromes, length 2 and up
    len2 = sum(1 for i in range(len(s) - 1) if s[i] == s[i + 1])
    return odd + even - len2

Problem 9: palindrome queries

Statement. Given s, then q queries l r (0-based, both included), print YES if s[l..r] is a palindrome, otherwise NO.

Input. Line 1: s. Line 2: q. Then q lines, each with l and r.

Output. q lines, each YES or NO.

Limits. 1 ≤ n ≤ 100000, 1 ≤ q ≤ 100000, 0 ≤ l ≤ r < n.

Example.

input
abacaba
3
0 6
1 3
2 4
output
YES
NO
YES

Hint. A query needs to compare a window with its reverse. Think about what you can compare in constant time after a linear preparation.

Target. O(n + q).

Show the answer

Build the hash of s and the hash of the reversed string. The window s[l..r] reads the same backwards as the window of the reversed string that starts at n - 1 - r. Two equal hashes mean YES. The reference is checked on 1000 random strings of length 12 with 6 random queries each (with l < r) over a and b, against a brute force that compares the window with its reverse. Sample: the example gives YES, NO, YES.

def p9_fast(s, queries):
    h, pw = hash_table(s)
    rs = s[::-1]
    hr, pr = hash_table(rs)
    n = len(s)
    out = []
    for l, r in queries:
        fwd = sub_hash(h, pw, l, r + 1)
        rev = sub_hash(hr, pr, n - 1 - r, n - l)
        out.append("YES" if fwd == rev else "NO")
    return out

Problem 10: weighted occurrences of a dictionary

Statement. You get m distinct words, each with a weight. Print the sum, over all occurrences of all words in the text, of the weight of the word. Occurrences may overlap, and one position may hold the end of several words.

Input. Line 1: m. Next m lines: a word and its weight. Then one line with the text t.

Output. One number.

Limits. 1 ≤ m ≤ 1000, the total length of the words is at most 100000, 1 ≤ weight ≤ 10^9, 1 ≤ |t| ≤ 10^6. Lowercase letters only. Use 64-bit integers.

Example.

input
3
a 1
ab 10
bab 100
abab
output
122

The word a occurs at positions 0 and 2, which gives 2. The word ab occurs at positions 0 and 2, which gives 20. The word bab occurs at position 1, which gives 100. The total is 122.

Hint. Think about a structure that reads the text once and knows, at each position, every word that ends there. One word can end inside another, so both must be counted.

Target. O(total word length + |t| + number of occurrences).

Show the answer

Build an Aho-Corasick automaton: a trie with failure links found by breadth-first search. Each node stores the total weight of the words that end at it, plus the total of its failure link, which covers words ending inside it. Walk the text once and add the stored total at each state. The reference is checked on 1000 random dictionaries of up to 4 words of length 1 to 3, with texts of length 1 to 14 over a and b, against a brute force that checks each word at each position. Sample: the example gives 122.

from collections import deque
 
 
def p10_fast(words, text):
    goto = [{}]
    fail = [0]
    val = [0]
    for w, wt in words:
        v = 0
        for ch in w:
            if ch not in goto[v]:
                goto.append({})
                fail.append(0)
                val.append(0)
                goto[v][ch] = len(goto) - 1
            v = goto[v][ch]
        val[v] += wt
    queue = deque()
    for ch, v in goto[0].items():
        fail[v] = 0
        queue.append(v)
    while queue:
        u = queue.popleft()
        for ch, v in goto[u].items():
            f = fail[u]
            while f and ch not in goto[f]:
                f = fail[f]
            fail[v] = goto[f][ch] if ch in goto[f] else 0
            val[v] += val[fail[v]]
            queue.append(v)
    total = 0
    state = 0
    for ch in text:
        while state and ch not in goto[state]:
            state = fail[state]
        if ch in goto[state]:
            state = goto[state][ch]
        total += val[state]
    return total

Problem 11: distinct palindromic substrings

Statement. Count the different palindromic substrings of s. A substring counts once however many times it occurs. Single letters count, since they are palindromes.

Input. One line with s.

Output. One number.

Limits. 1 ≤ n ≤ 2000. Lowercase letters only.

Example.

input
aabaa
output
5

The distinct palindromes are a, b, aa, aba and aabaa.

Hint. Two equal palindromes at different places must be counted once. Find all palindromes first, then think about how to tell equal ones apart.

Target. O(n^2) with n ≤ 2000.

Show the answer

Use the Manacher radii to find every palindrome: each centre contributes a nested family of palindromes, one for each radius. Store the hash of each palindrome in a set, and the answer is the size of the set. The total number of palindromes can be O(n^2), so this takes O(n^2) time at the limit. The reference is checked on 2000 random strings of length 1 to 14 over a and b, against a brute force that tests every substring and puts the palindromes in a set. Sample: aabaa gives 5.

def p11_fast(s):
    p = manacher_radii(s)
    n = len(s)
    h, pw = hash_table(s)
    seen = set()
    for i in range(1, len(p)):
        if i % 2 == 1:
            c = (i - 1) // 2  # odd centre s[c]
            for k in range(0, p[i] // 2 + 1):
                seen.add(sub_hash(h, pw, c - k, c + k + 1))
        else:
            c = i // 2  # even centre between s[c - 1] and s[c]
            for k in range(1, (p[i] + 1) // 2 + 1):
                seen.add(sub_hash(h, pw, c - k, c + k))
    return len(seen)

Problem 12: the smallest rotation

Statement. A rotation of s moves a prefix of s to its end. Print the 1-based position of the rotation that is smallest in dictionary order. If two rotations are equal, print the smaller position.

Input. One line with s.

Output. One number.

Limits. 1 ≤ n ≤ 100000. Lowercase letters only.

Example.

input
bca
output
3

The rotations of bca start at positions 1, 2 and 3 (1-based): bca, cab, abc. The smallest is abc, which starts at position 3.

A second example: for abab the rotations starting at positions 1 and 3 are equal (abab), so the answer is 1.

Hint. Think about the string written twice, and about a scan that keeps the best start so far and skips starts that cannot win.

Target. O(n).

Show the answer

Use the minimum-rotation scan, which keeps two candidate starts and skips the blocks that cannot hold the smallest rotation. It works on s + s. The reference below returns the 0-based start of the first smallest rotation, and the answer adds 1. It is checked on 3000 random strings of length 1 to 12 over a and b, against a brute force that sorts all rotations with their positions. Samples: bca gives 3, abab gives 1.

def min_rotation_start(s):
    s2 = s + s
    n = len(s)
    i = 0
    ans = 0
    while i < n:
        ans = i
        j, k = i + 1, i
        while j < 2 * n and s2[k] <= s2[j]:
            if s2[k] < s2[j]:
                k = i
            else:
                k += 1
            j += 1
        while i <= k:
            i += j - k
    return ans
 
 
def p12_fast(s):
    return min_rotation_start(s) + 1

Problem 13: fewest insertions to make a palindrome

Statement. You may insert letters anywhere in s. Print the fewest insertions that make the result a palindrome.

Input. One line with s.

Output. One number.

Limits. 1 ≤ n ≤ 2000. Lowercase letters only.

Example.

input
abca
output
1

Insert one b between the c and the last a to get abcba. The second example is abc, which needs 2 insertions, for example to get abcba.

Hint. The letters you do not insert already form a subsequence. Ask which subsequence is the longest palindrome, and what that means for the other letters.

Target. O(n^2).

Show the answer

The letters that stay put form a palindromic subsequence, and every other letter needs one insertion. So the answer is n minus the length of the longest palindromic subsequence. That length is the longest common subsequence of s and its reverse, which a two-row dynamic program computes. The reference is checked on 300 random strings of length 1 to 5 over a, b and c, against a brute force that tries every palindrome over the letters of s in order of length and checks that s is a subsequence of it. Samples: abca gives 1, abc gives 2.

def p13_fast(s):
    r = s[::-1]
    n = len(s)
    prev = [0] * (n + 1)
    for i in range(1, n + 1):
        cur = [0] * (n + 1)
        for j in range(1, n + 1):
            if s[i - 1] == r[j - 1]:
                cur[j] = prev[j - 1] + 1
            else:
                cur[j] = max(prev[j], cur[j - 1])
        prev = cur
    return n - prev[n]

Problem 14: matching with wildcards

Statement. The pattern p may contain the letter ?, which matches any single letter of the text. Count the positions i where p matches t starting at i.

Input. Line 1: t. Line 2: p.

Output. One number.

Limits. 1 ≤ |p| ≤ |t| ≤ 2000. The text has lowercase letters. The pattern has lowercase letters and ?.

Example.

input
abcabc
a?c
output
2

The pattern matches at positions 0 and 3: abc and abc.

Hint. Why does the prefix function stop working when the pattern can match more than one letter? Think about what a border means when a ? can stand for any letter.

Target. O(|p| · |t|) is fine at this limit. Faster methods based on convolution exist for large inputs, but this problem does not need them.

Show the answer

Compare the pattern with every window of the text, letter by letter, and stop at the first mismatch. A ? matches anything. The prefix function needs equality to be transitive, and ? breaks that: a equals ?, and ? equals b, but a does not equal b. The reference is checked on 2000 random pairs with t of length 1 to 14 and p of length 1 to 4 over a, b and ? in the pattern, against a brute force that uses a regular expression in which each ? becomes any letter. Sample: abcabc, a?c gives 2.

def p14_fast(t, p):
    n, m = len(t), len(p)
    out = 0
    for i in range(n - m + 1):
        ok = True
        for j in range(m):
            if p[j] != "?" and p[j] != t[i + j]:
                ok = False
                break
        out += ok
    return out

Common mistakes

A brute force that shares a bug with the fast code. If both use the same hash helper with the same error, the random test passes and the answer is still wrong. Write the brute force without the helper. Use plain slices or a set of strings.

Ties and empty ranges. Problem 12 has ties, and problem 9 can ask for a window of length 1. Test these by hand. The random tests above include them when the alphabet is small.

Counting a palindrome twice. In problem 11, the same palindrome appears at many centres. Put the strings in a set. Counting the nested palindromes at each centre without a set gives a number that is too large.

Leaving out the word that ends inside another word. In problem 10, a text that contains bab also contains ab at the end of the word, and ab must be counted too. A walk that only reads the state at the end of each word misses it.

Wrong count on the sample. If your answer for problem 2 on aaaa is 4, you counted the single letter a. The statement asks for exponent at least 2.

Checklist

• Write the brute force first, and keep it free of the fast code's helpers.
• For every answer, test a case with one letter, a case where the pattern is as long as the text, and a case with ties.
• Use 64-bit integers for every count, sum and weight.
• Keep a random seed in your test script, so a failing case can be run again.
• When a hash decides equality, use two moduli or a random base, and check a collision by comparing the strings.
• Time the fast solution at the largest limit, and compare with the target in the statement.

Summary

The set mixes the module's tools on purpose, so the first step on each problem is the choice of tool: the prefix function for periods and overlaps, hashing for equal windows and for palindrome tests, Manacher for palindrome centres, automata and tries for many strings at once, Duval-style scans for rotations, and dynamic programs when the answer is a best cut or a best alignment. Your own brute force is the judge of every answer. The Python references above are checked against brute forces on small random inputs, and the C++ versions are yours to write and time.