Prefix function and KMP

Why this matters when you build with AI

When you ask an AI tool for a function that finds a word in a long file, it often returns two nested loops: for every start position, compare the whole pattern. That code is correct, and it is fast on short inputs. On a text of 200000 characters with a pattern of 100000 characters, the comparisons can add up to billions of steps, and a judge with a time limit gives a "time limit exceeded" result. The tool does not run the code on a large input, so it does not see the problem.

This lesson teaches a method that does the same job in time proportional to the length of the text plus the length of the pattern. The main tool is the prefix function. You will build it, prove its running time, and use it for two tasks: counting occurrences and listing borders. You will also see the same table used as a small machine that reads the text one character at a time.

To catch an AI mistake in this area, do two things. First, read the loops and count how many times each index moves. Second, run the code on a random small input and compare it with a brute force program, which is a program that tries every possibility with no clever idea. Both checks take a few minutes and catch the common errors listed later in this lesson.

Terms used in this lesson

A string is a sequence of characters. A prefix of length k is its first k characters. A suffix of length k is its last k characters. A prefix or suffix is proper when it is not the whole string. A border is a string that is both a proper prefix and a proper suffix of the same string. The string aba is a border of abacaba, because it starts the string and ends it.

The pattern is the string you search for. The text is the string you search in. An occurrence is a start position where the pattern matches the text. Occurrences may overlap, so aa occurs three times in aaaa (at positions 0, 1 and 2).

An amortised bound is an average over all steps. One step can be expensive, but the total over all steps is still small. This matters for the running time proof below.

The idea: the longest border of each prefix

Take a string s with n characters, indexed from 0. For each index i, look at the prefix s[0..i], which has length i + 1. The prefix function pi[i] is the length of the longest proper border of that prefix. If there is no border, the value is 0.

A worked example makes this concrete. Take s = abacaba.

The last row has two borders, a and aba. The prefix function stores only the longest one, so pi[6] is 3. The array is 0 0 1 0 1 2 3. This is the output of the program in the code section below.

Two facts make the rest of the lesson work.

1. Every value satisfies pi[i] <= i, because a proper border of a prefix of length i + 1 is at most i characters long.
2. If pi[i] is k, then pi[i-1] is at least k - 1. The border of length k loses its last character and becomes a border of s[0..i-1].

Fact 2 gives the method. To compute pi[i], start from the candidate length k = pi[i-1]. Extend it by one character if s[i] equals s[k]. If not, the next candidate is the border of the border, which is pi[k-1]. Repeat until a match is found or the candidate is 0.

Building the table in linear time

The code keeps one number, k, which is the length of the current border. Each step either extends k by one or lowers it. The next section proves that the total work is linear.

vector<int> prefix_function(const string& s) {
    int n = s.size();
    vector<int> pi(n, 0);
    for (int i = 1; i < n; i++) {
        int k = pi[i - 1];
        while (k > 0 && s[i] != s[k]) k = pi[k - 1];
        if (s[i] == s[k]) k++;
        pi[i] = k;
    }
    return pi;
}

Notice that pi[0] is always 0, so the loop starts at 1. The check k > 0 must come first in the while, because pi[k - 1] with k = 0 reads the index -1.

Why the total work is O(n)

Look at the variable k over the whole run. It starts at 0. In each iteration of the outer loop, k can grow by at most 1, because k starts at pi[i-1], and pi[i] <= pi[i-1] + 1. The while loop only lowers k, and each pass lowers it by at least 1, because pi[k-1] < k.

Let the total number of increases be at most n, one per outer iteration. The number of decreases can never be more than the number of increases, because k never goes below 0. So the total number of passes through the while loop is at most n. Each outer iteration does a constant amount of work plus its passes, so the whole function is O(n). The argument is amortised: one outer step may do many passes, but the passes are paid for by earlier increases.

Worked problem: counting code words in a gate log

Here is an original problem to solve by hand first.

A gate system writes all swipe codes into one long log, one lowercase string. A code word is a lowercase string. Count how many times the code word appears in the log. Overlapping occurrences count separately.

Input: line 1 holds the code word p, line 2 holds the log t. Both use lowercase letters only. Limits: 1 <= |p| <= 200000 and 1 <= |t| <= 200000. Output: one integer, the number of occurrences.

Example: the code word ana and the log bananaanana give 4. The occurrences start at positions 1, 3, 6 and 8.

A direct search costs |p| times |t| in the worst case. The prefix function gives a linear method. Build the string p + '#' + t, where # is a character that does not appear in p or t. Compute the prefix function of this combined string. An occurrence of p ends at a position i in the text part exactly when pi[i] equals |p|.

Why does the separator work? The # appears only once in the combined string, at index |p|. Take any index i after the separator. A border of the prefix ending at i that is longer than |p| would contain the #, and the copy of that border at the end of the prefix would need a # at the same offset. The only # in the string is at index |p|, so the two copies are the same position, which means the border is not proper. So pi[i] never goes above |p| in the text part, and the value |p| marks a full match exactly.

The complete program for this problem follows. It also handles the border task from the practice section, so the same file serves both.

#include <bits/stdc++.h>
using namespace std;
 
// pi[i] = length of the longest proper border of s[0..i].
vector<int> prefix_function(const string& s) {
    int n = s.size();
    vector<int> pi(n, 0);
    for (int i = 1; i < n; i++) {
        int k = pi[i - 1];
        while (k > 0 && s[i] != s[k]) k = pi[k - 1];
        if (s[i] == s[k]) k++;
        pi[i] = k;
    }
    return pi;
}
 
// Number of occurrences of p in t, overlaps included.
long long count_occurrences(const string& p, const string& t) {
    int m = p.size();
    string s = p + '#' + t;  // '#' must not appear in p or t
    vector<int> pi = prefix_function(s);
    long long cnt = 0;
    for (int i = m + 1; i < (int)s.size(); i++)
        if (pi[i] == m) cnt++;
    return cnt;
}
 
// Lengths of all proper borders of s, in increasing order.
vector<int> all_borders(const string& s) {
    vector<int> pi = prefix_function(s);
    vector<int> res;
    int k = pi.empty() ? 0 : pi.back();
    while (k > 0) {
        res.push_back(k);
        k = pi[k - 1];
    }
    reverse(res.begin(), res.end());
    return res;
}
 
int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    string task;
    cin >> task;
    if (task == "count") {
        string p, t;
        cin >> p >> t;
        cout << count_occurrences(p, t) << "\n";
    } else {
        string s;
        cin >> s;
        vector<int> b = all_borders(s);
        for (size_t i = 0; i < b.size(); i++) cout << b[i] << (i + 1 < b.size() ? " " : "\n");
        if (b.empty()) cout << "\n";
    }
}

Compile and run it with g++ 15.2.0 (the version on the machine where this lesson was tested) using g++ -std=c++17 -O2 -o kmp kmp.cpp. The output on the example is:

$ printf 'count\nana\nbananaanana\n' | ./kmp
4

To see the table inside the combined string, run the prefix function on ana#bananaanana. The output is:

0 0 1 0 0 1 2 3 2 3 1 2 3 2 3

The indices with value 3 are 7, 9, 12 and 14. Subtract |p| + 1 = 4 to get the text positions 3, 5, 8 and 10. These are the ends of the occurrences. The start of each occurrence is the end minus |p| + 1, which gives 1, 3, 6 and 8, the same as in the example statement.

Testing the solution

I compared count_occurrences with a brute force loop that checks every start position. The test ran 500 random cases over the alphabet abc and ab, with patterns of 1 to 4 characters and texts of 0 to 14 characters. The two programs gave the same answer on every case. The same test also compared all_borders with a brute force that tests every length, on 500 random strings of length 1 to 14, with no differences. The random generator and the brute force are not part of this lesson's code, but the steps are simple enough to repeat in your own folder.

The automaton view

The prefix function can also drive a machine. Let state k mean "the last k characters read match the first k characters of the pattern". State 0 is the start, and state m (the pattern length) means a full match. From state k, reading a character c has two cases. If p[k] equals c, the machine moves to state k + 1. If not, it falls back to the state of the longest border, and tries c from there. That fallback state is pi[k-1].

Precomputing the transitions for every state and every letter gives a table with one row per state. Then each character of the text costs one table lookup. The table itself is built with the prefix function, so the total time is linear in the pattern length times the alphabet size.

// aut[k][c] = next state after reading letter c in state k, for k = 0..m.
vector<array<int, 26>> aut(m + 1);
for (int k = 0; k <= m; k++)
    for (int c = 0; c < 26; c++) {
        if (k < m && p[k] - 'a' == c) aut[k][c] = k + 1;
        else if (k == 0) aut[k][c] = 0;
        else aut[k][c] = aut[pi[k - 1]][c];
    }

The row for state k depends only on rows with smaller index, because pi[k-1] < k, so the loop can fill the table in order. I tested a complete program that builds this table and counts occurrences in the text. It matched the brute force on 400 random cases, with no differences. On the input aba and ababa it prints 2, the number of occurrences at positions 0 and 2.

Use the automaton when the text is read as a stream and you cannot store it, or when you need to run many texts against one pattern. Use the plain prefix function when you need the whole table, for borders or periods.

Complexity

Building the prefix function takes O(n) time and O(n) memory, where n is the length of the string. Counting occurrences with the separator takes O(|p| + |t|) time and memory. The automaton takes O(|p| times the alphabet size) memory and O(|t|) time after the table is built. The running times are amortised for the prefix function, and they are worst-case bounds for the other parts. Constant factors matter in the real contest, so measure your code on a large random input before you trust it.

Common mistakes

1. No separator between pattern and text. The pattern aa in the text a has no occurrence. Without a separator, the code joins them as aaa, and the prefix function reports a full match at index 2. A variant of the program with no separator printed 1, while the correct answer is 0. The version with # between the two strings printed 0. The fix is to insert a character that is not in either string.

2. Counting border positions as occurrences. The value pi[i] describes the prefix ending at i, not the pattern. Only the positions where pi[i] equals |p| are occurrences. A common error is to count all positions where pi[i] is positive, which counts partial matches.

3. Missing k > 0 in the while loop. When k is 0, the expression pi[k - 1] reads index -1. This is undefined behaviour in C++. The result depends on the compiler and the input, so the bug may hide on small tests and show on large ones. Keep the check in the while condition.

4. Reporting the end position instead of the start. The loop gives the index where a match ends. The example above gives ends 3, 5, 8 and 10 for the pattern ana in bananaanana, and starts 1, 3, 6 and 8. A judge that asks for start positions gives a wrong answer if you print the ends.

5. Confusing a border with a substring. A border must be both a prefix and a suffix. The string bac occurs inside abacaba, but it is not a border, because it does not start the string. Check both ends before you add a length to the list.

Hands-on exercise

Work in a new folder on your machine.

1. Save the program from the worked problem as kmp.cpp and compile it with g++ -std=c++17 -O2 -o kmp kmp.cpp.
2. Write a file brute.py that reads the same input and counts occurrences with a loop over all start positions. Use text.startswith(pattern, i) in the loop.
3. Write a generator gen.py that prints a random pattern and a random text, each made of the letters a and b, with lengths from 1 to 5 and 0 to 20.
4. Run both programs on 300 random inputs with a shell loop. Stop and report any input where the outputs differ.
5. Then change the pattern to a string with length 0 and observe what your code does. Decide what the right answer is and fix the code if needed.

What a good result looks like: all 300 outputs match, and you can explain the answer for the empty pattern in one sentence. If the outputs differ, print the input that failed and check it by hand with the table method.

Prompt for checking an AI answer

Use this prompt when an AI tool gives you a pattern search function. It asks the tool to justify its own code instead of just pasting it.

Here is my function count_occurrences(p, t). Check it in three steps.
1. Explain, for the string p + '#' + t, which indices count as matches and why.
2. Give the smallest input where a missing separator changes the answer, and the two results.
3. State the running time in terms of |p| and |t|, and say whether the bound is worst case or amortised.
Do not rewrite the code. Only report errors and their smallest inputs.

A good answer names the indices where the value equals |p|, gives an input such as pattern aa and text a, and states O(|p| + |t|) time. Check the input yourself with the program. An answer that says the code is "optimal" without giving a proof is not enough.

Checklist

• Build the combined string with a character that appears in neither input.
• Check k > 0 before reading pi[k - 1].
• Count positions where the value equals the pattern length, not any positive value.
• Convert end positions to start positions before printing.
• Test the function against a brute force on at least a few hundred random small inputs.
• Measure the running time on a large random input before you trust the bound.

Summary

The prefix function gives, for each prefix of a string, the length of its longest proper border. A single number k moves up by at most one per character, and every lowering step is paid for by an earlier raise, so the table takes linear time. Put a separator between the pattern and the text to count occurrences in linear time, and read the same table as an automaton when the text is streamed.

Practice problems

Practice 1: All borders of a string

Write a program that reads one string s of lowercase letters, with length from 1 to 200000, and prints the lengths of all its proper borders in increasing order. If there are none, print an empty line.

Show the answer

Hint: the longest border is pi[n-1]. Every other border is a border of that border, so follow the chain pi[k-1] from the longest one down to 0 and reverse the list. The code is all_borders in the worked problem. On the input abacaba it prints 1 3. On the input banana it prints an empty line.

Practice 2: Smallest period

A string s has period p when s[i] == s[i+p] for every i where both indices exist. The smallest period of a string of length n is its shortest such p. Write a program that reads a string and prints its smallest period.

Show the answer

Hint: a border of length L gives the period n - L. Take the longest border, which is pi[n-1], so the answer is n - pi[n-1]. On abacaba the longest border is 3, so the period is 7 - 3 = 4. You can check that 4 is the smallest by hand: periods 1, 2 and 3 each fail on some index. On aabaaab the longest border is 3, so the period is 4.

Practice 3: The table by hand

Compute the prefix function of aabaaab by hand, one index at a time. Write down the value of k at each step. Then run the program and compare.

Show the answer

Hint: for each index i, start from the value of pi[i-1], extend it when the characters match, and fall back through the chain otherwise. The table is 0 1 0 1 2 2 3. Index 4 extends the border a to aa, giving 2. Index 5 fails to extend aa because s[5] is a and s[2] is b, so it falls back to pi[1], which is 1, and then s[5] equals s[1], so the value becomes 2. Index 6 extends 2 to 3, because s[6] is b and s[2] is b.

Practice 4: Start positions

Count the occurrences of aba in ababababa and print their start positions.

Show the answer

Hint: the count is the number of end positions, and each start is the end minus the pattern length plus 1. The occurrences start at 0, 2, 4 and 6, so the count is 4. The program count_occurrences gives 4. To print starts, collect the indices i where the value equals the pattern length, then subtract the pattern length minus 1 from each.

Practice 5: The automaton by hand

Build the automaton for the pattern aab on the letters a and b. Trace the states while reading the text aaab. Say how many times the state reaches 3.

Show the answer

Hint: state 0 reads a to state 1, then a to state 2. From state 2, reading a fails because the pattern has b there, so the machine falls back to state pi[1], which is 1, and a moves it to state 2. Reading b moves to state 3. The state sequence is 1, 2, 2, 3. The state reaches 3 once, so the answer is 1, and the program aut gives 1 on this input.