Prefix function

The prefix function computes, for every position of a string, the length of the longest stretch ending there that repeats the start of the string, in $O(n)$ time. With it you can find every occurrence of a pattern in a text read one letter at a time (the KMP algorithm), or the shortest piece a string repeats (its period). The Z-function solves the same problems with a different array.

Explanation

Where does the pattern ab occur in the text aabab? Glue the pattern, a # and the text into one string, ab#aabab, 8 letters at positions 0 to 7. Under every position, write the length of the longest stretch that ends there and repeats the start of the string. The stretch may not begin at position 0 itself, so position 0 gets 0. Position 5 gets 2: positions 4 and 5 hold ab, like the start of the string, while 3 letters ending there, aab, don't repeat it. An occurrence of ab that ends at some position of the text is exactly a stretch of length 2 there, so finding the occurrences means computing these numbers.

Fig. 1. Positions 5 and 7 get 2.

The 2s stand at positions 5 and 7, so ab starts 2 letters earlier, at positions 4 and 6. These are positions 1 and 3 of aabab.

A program could find these numbers the plain way: at each position, try every length from the longest down, and compare letter by letter from the left until a pair differs or the length fits. On ab#aabab that is cheap, 36 comparisons. The cost shows on strings that repeat themselves: aabaabaaabb (positions 0 to 10) takes 92 comparisons for its 11 numbers. For 10^6 letters a it takes about 5 \cdot 10^{11}, far beyond the rule of thumb of about 10^8 simple steps per second.

Take a stretch ending at position 8 that repeats the start, and drop its last letter: what is left is a stretch ending at position 7 that repeats the start. So each number can begin from the one before it: try to make the previous stretch one letter longer. Where the letter doesn't fit, as at position 2 (b after a), compare with the first letter instead: b is not a, so 0. For aabaabaaabb this gives 0 1 0 1 2 3 4 5 at positions 0 to 7. At position 7, the 5 letters at positions 3 to 7, aabaa, repeat positions 0 to 4.

Fig. 2. Position 8 can't extend the stretch.

Position 8 holds a, but the start of the string continues after its first 5 letters with position 5, b. The 5-letter stretch can't be extended.

The long stretch is broken, so the natural move is to start over, as at position 2: write 1 if position 8 holds the same letter as position 0, and 0 otherwise. Position 8 holds a, so this gives 1.

Fig. 3. Starting over writes 1; the true value is 2.

Positions 7 and 8 hold aa, and the string starts with aa too. Starting over threw away the shorter stretches ending at position 7, and one of them, the single a at position 7, could still be extended. The error spreads: position 9 then gets 0 instead of 3.

Which shorter stretches end at position 7 and repeat the start? Since positions 3 to 7 equal positions 0 to 4, a stretch at the end of one that repeats the start is also one at the end of the other. The longest one is the number at position 4, already written: 2. Try to extend it: the start continues with position 2, b, and position 8 is a, so no. The next shorter one comes from the same reasoning on those 2 letters: the number at position 1, which is 1. Try to extend it: the start continues with position 1, a, and position 8 is a, so yes. The number at position 8 is 1 + 1 = 2, the true value.

Fig. 4. Each shorter stretch is the number at the end of the one before it. Only the third one can be extended.

This row of numbers is the prefix function of the string. It is filled left to right: at each position, try to extend the previous stretch by one letter, and while the letter doesn't fit, fall back to the next shorter stretch, whose length is the number written at the end of the current one. When no stretch is left, compare with the first letter of the string: write 1 if it matches and 0 if not.

Here is the whole run on aabaabaaabb, from position 1 to 10.

Position 1 holds a, the same as position 0, so the number at position 1 is 1. The pointer i marks the position being computed.

Position 2: take the stretch ending at position 1, 1 letter (a); it would be extended by position 1, a.

Position 2 holds b and position 1 holds a: they differ, so the stretch can't be extended.

Fall back: the copy at the start ends at position 0, whose number is 0, so no stretch is left and position 2 is compared with position 0.

Position 2 holds b and position 0 holds a: they differ, so the number at position 2 is 0.

The number at position 2 is 0, so position 3 (a) is compared with position 0 (a): they are equal, so the number is 1.

Extend: the stretch ending at position 3 has 1 letter, and position 4 (a) equals position 1, the letter after its copy at the start, so the number is 2.

Extend: the stretch ending at position 4 has 2 letters, and position 5 (b) equals position 2, the letter after its copy at the start, so the number is 3.

Extend: the stretch ending at position 5 has 3 letters, and position 6 (a) equals position 3, the letter after its copy at the start, so the number is 4.

Extend: the stretch ending at position 6 has 4 letters, and position 7 (a) equals position 4, the letter after its copy at the start, so the number is 5.

Position 8: take the stretch ending at position 7, 5 letters (aabaa); it would be extended by position 5, b.

Position 8 holds a and position 5 holds b: they differ, so the stretch can't be extended.

Fall back: the copy at the start ends at position 4, whose number is 2, so the next shorter stretch has 2 letters (aa).

Position 8 holds a and position 2 holds b: they differ, so the stretch can't be extended.

Fall back: the copy at the start ends at position 1, whose number is 1, so the next shorter stretch has 1 letter (a).

Extend: position 8 holds a, like position 1, so the stretch grows to 2 letters and the number at position 8 is 2.

Extend: the stretch ending at position 8 has 2 letters, and position 9 (b) equals position 2, the letter after its copy at the start, so the number is 3.

Position 10: take the stretch ending at position 9, 3 letters (aab); it would be extended by position 3, a.

Position 10 holds b and position 3 holds a: they differ, so the stretch can't be extended.

Fall back: the copy at the start ends at position 2, whose number is 0, so no stretch is left and position 10 is compared with position 0.

Position 10 holds b and position 0 holds a: they differ, so the number at position 10 is 0.

The whole row is filled: 0 1 0 1 2 3 4 5 2 3 0, with 14 comparisons and 4 fallbacks.

1/22

Fig. 5. Step through the prefix function with the buttons or the arrow keys. The blue brackets are the current stretch and the start it repeats; green pairs are equal, red pairs differ.

The run fills in the whole array, 0 1 0 1 2 3 4 5 2 3 0, with 14 comparisons instead of 92, falling back 4 times. For 10^6 letters a every stretch extends at once, and the run makes 10^6 - 1 comparisons instead of about 5 \cdot 10^{11}.

Searching for a pattern

To find a pattern of length $m$ in a text of length $n$, glue the pattern, a separator and the text, and compute the prefix function. A number equal to $m$ at glued position $i$ means the whole pattern ends there, so an occurrence starts at text position i - 2m (in ab#aabab, position 5 gives 5 - 4 = 1). The separator must appear in neither string; then no stretch runs through it, and no number is larger than $m$. Without a separator, the pattern aa and the text aaa give aaaaa with the numbers 0 1 2 3 4. The stretches run across the join, the two occurrences get 3 and 4 instead of 2, and the test "equals $m$" finds neither.

A number after the # needs only the letter at its position, the number just before it, and, for falling back, numbers at positions below $m$, which belong to the pattern. Nothing else from the text is used. So we can compute the prefix function of the pattern once, keep one number for the current stretch, and read the text letter by letter without storing it. This is the Knuth–Morris–Pratt algorithm, KMP. After an occurrence, the current stretch drops to the pattern's last number, so occurrences that overlap are found too. That drop also keeps the number at most $m$, so KMP needs no separator.

Borders and periods

The stretches the algorithm extends are called borders: a border of a string is a stretch, shorter than the whole string, that is both its start and its end. The number at position $i$ is the length of the longest border of the first i + 1 letters. Take aabaabaa, the first 8 letters of the main string (numbers 0 1 0 1 2 3 4 5). Its longest border is the last number, 5: aabaa. The next shorter border comes from the same fallback as in the algorithm: the number at position 4, 2, gives aa, the number at position 1, 1, gives a, and the number at position 0 ends the list.

A period of a string is a shift $p$ such that every letter equals the one $p$ places before it. A border of length 5 means that the string, shifted right by 8 − 5 = 3, matches itself where the copies overlap.

Fig. 6. Every letter equals the one 3 places before it.

So 3 is a period. In general a period $p$ is the same thing as a border of length 8 - p, so the longest border gives the shortest period, 3.

Since 8 is not divisible by 3, the string is not aab written several times; the shortest such piece is the whole string. The same test works on every start of the string: the first 6 letters, aabaab, have the number 3 at position 5, and 6 − 3 = 3 divides 6, so aabaab is aab twice.

Implementation

We keep the numbers in a vector pi of size n. The start of a string is called its prefix, so pi[i] is the length of the longest border of the prefix ending at position i, and pi[0] = 0. (Other sources call this array the failure function.)

While computing pi[i], the variable k is the length of the current border of the prefix ending at position i - 1. We try to extend it with s[i]. The letter that would extend it is s[k], the first one after the prefix of length k. The border ends at position k - 1 of that prefix, so the next shorter border has length pi[k - 1].

For each position i from 1 to n - 1:

1. Take the previous border: k = pi[i - 1].

2. Fall back: while k > 0 and s[i] != s[k], set k = pi[k - 1].

3. Extend: if s[i] == s[k], increase k. Then pi[i] = k.

At i = 8, k takes the values 5, 2 and 1, the three stretches tried at position 8 above, then extends to 2.

The program reads one line, the string, and prints its prefix function. For the main example the input is:

aabaabaaabb

#include <iostream>
#include <string>
#include <vector>
using namespace std;

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];                              // take the previous border
        while (k > 0 && s[i] != s[k]) k = pi[k - 1];    // fall back
        if (s[i] == s[k]) k++;                          // extend
        pi[i] = k;
    }
    return pi;
}

int main() {
    string s;
    cin >> s;
    vector<int> pi = prefix_function(s);
    int n = pi.size();
    for (int i = 0; i < n; i++) cout << pi[i] << (i + 1 < n ? ' ' : '\n');
}

The prefix function runs in $O(n)$ time (the proof is below) and uses $O(n)$ memory for s and pi.

Searching with KMP

The program reads the text t, then the pattern p, and prints how many times p occurs in t and where each occurrence starts, counting from 0. Here k is the length of the longest start of p that ends at the last letter read. It carries over from one letter of the text to the next, so step 1 has nothing to do; the other two steps are the same, with p[k] as the letter that would extend it.

#include <iostream>
#include <string>
#include <vector>
using namespace std;

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];                              // take the previous border
        while (k > 0 && s[i] != s[k]) k = pi[k - 1];    // fall back
        if (s[i] == s[k]) k++;                          // extend
        pi[i] = k;
    }
    return pi;
}

int main() {
    string t, p;
    cin >> t >> p;
    int m = p.size();
    vector<int> pi = prefix_function(p);   // only the pattern's array is kept

    vector<int> found;
    int k = 0;   // length of the longest start of p ending at the last letter read
    for (int i = 0; i < (int)t.size(); i++) {
        while (k > 0 && t[i] != p[k]) k = pi[k - 1];   // fall back
        if (t[i] == p[k]) k++;                         // extend
        if (k == m) {                                  // an occurrence ends at i
            found.push_back(i - m + 1);                // position in t, from 0
            k = pi[m - 1];
        }
    }

    cout << found.size() << '\n';
    for (int j = 0; j < (int)found.size(); j++)
        cout << found[j] << (j + 1 < (int)found.size() ? ' ' : '\n');
}

For the text aabab and the pattern ab it prints 2 occurrences, at positions 1 and 3. To avoid storing the text, replace the loop over t with a loop that reads one letter at a time with getchar(); the body stays the same. When a problem numbers positions from 1, push i - m + 2 instead.

Correctness proof

Write \pi_i for the number at position $i$, and call a border of the prefix ending at position $i$ a border at $i$.

First, the borders at $i$, longest first, are \pi_i, then \pi_{\pi_i - 1}, then the number at the end of that one, and so on down to 0. Let a = \pi_i be the longest border, and let b < a. The border of length $a$ is both the start and the end of the prefix. So its first $b$ letters are the first $b$ letters of the prefix, and its last $b$ letters are the last $b$ letters of the prefix. Hence a stretch of length $b$ is a border at $i$ exactly when it is a border of the border of length $a$, that is, a border at position a - 1. The longest such stretch is \pi_{a - 1}, and repeating the argument lists all the others in order.

Fig. 7. Each shorter border is a border of the longest one.

Second, every non-empty border at $i$, without its last letter, is a border at i - 1 (possibly empty), and a border of length $j$ at i - 1 extends by s_i exactly when s_i = s_j. The algorithm keeps one invariant: before each comparison, $k$ is a border at i - 1, and no longer border at i - 1 extends. It starts with the longest, k = \pi_{i-1}, and each fallback moves to the next shorter border, by the first part. So it stops at the longest border that extends and writes \pi_i = k + 1; if even the empty border doesn't extend (s_i \ne s_0), it writes 0.

The running time follows from the value of $k$. It grows by at most 1 per position, so by at most n - 1 in total. Each fallback lowers it by at least 1, and it never goes below 0, so there are fewer than $n$ fallbacks over the whole run. Each fallback costs one comparison, and each position ends with one more, so the run makes fewer than $2n$ comparisons. For aabaabaaabb that is 10 positions plus 4 fallbacks, the 14 comparisons of the run above.

Common mistakes

• Falling back with if instead of while. For aaab the last number becomes 1 instead of 0: after one fallback k = 1, and s[3] = b still differs from s[1] = a, but the if doesn't fall back again, so the stale 1 is written. Repeat the fallback until the letter fits or k is 0.

• In KMP, setting k = 0 after an occurrence. This misses overlapping occurrences: the pattern aba in ababa is found only at position 0, not at 2. Set k = pi[m - 1]; leaving k = m reads p[m], past the end of the pattern.

• Taking n - pi[n - 1] as the length of the piece that s repeats. For aabaabaa that is 8 − 5 = 3, but 8 is not divisible by 3, so the answer is 8, not aab. Check n % (n - pi[n - 1]) == 0.

Similar algorithms

The Z-function stores, for each start position, the length of the match with the prefix. Both run in $O(n)$. KMP reads the text one letter at a time and keeps only the pattern, while the Z-function needs the glued string in memory.

String hashing compares any two substrings in $O(1)$ after $O(n)$ preparation, with a small chance of error. Use it when the pieces to compare aren't tied to the start of one string, such as a common substring of a given length in two strings.

Aho–Corasick searches for many patterns at once: it builds the same fallback links over a tree of all the patterns.

Practice