Embercode
✕
TriesLesson 11 of 14

40 min read

Suffix trie for substring queries in Python

Why it matters

Prefix search asks whether a stored word starts with some letters. Substring search asks a harder question: does a pattern appear anywhere inside a text? Does ana appear in banana? It appears twice, at positions 1 and 3, and it does not start the word at all. A trie of whole words cannot answer this, because ana is not a stored word.

A suffix trie answers it with the same walk a prefix search uses. It stores every ending of the text, so any substring is the start of one of the stored endings. The price is memory. The trie grows much faster than the text, and this lesson measures how fast, so you can see where the method stops being practical.

Two definitions

A substring is any run of letters inside a text. It has a start position and an end position. The text banana has 15 distinct substrings that are not empty, such as b, an, nan and banana itself.

A suffix is a substring that runs to the end of the text. The text banana has six suffixes, one for each start position. They are banana, anana, nana, ana, na and a. In Python, the slice text[i:] gives the suffix that starts at position i.

The key fact is this. Every substring starts at some position i, and it is therefore a prefix of the suffix that starts at i. The substring nan starts at position 2, and the suffix starting at position 2 is nana, which begins with nan. So if you insert all suffixes into a trie, a substring is in the text exactly when the trie has a path that spells it. Checking a pattern is then a prefix walk.

Worked example one: banana by hand

Start with the six suffixes. Insert banana first. The path is b, a, n, a, n, a, and each new letter creates a new node, because the trie is empty.

Insert anana next. The root has no a child yet, so the path starts there and creates five new nodes. Insert nana. The root has an n child only if an earlier suffix started with n, and none did, so this path is new as well. Insert ana. The root has an a child from anana. The next letter, n, is also present under it. The next letter, a, is present too. So the path a, n, a already exists, and inserting ana only adds one to the counts of nodes that are already there. Insert na and a in the same way.

The walk does the same thing for every suffix. A node is created only when no earlier suffix passed through it. This is why the node count is not simply the total number of letters. Shared beginnings are stored once.

Each node in this lesson's code has a count, pass_count, that says how many suffixes pass through it. The count matters for queries, and it is easiest to see on the small example. The node for a at depth 1 is passed by the suffixes anana, ana and a, so its count is 3. The node for an is passed by anana and ana, so its count is 2. The node for nan is passed by nana only, so its count is 1.

These three counts are the numbers of times a, an and nan occur in banana. The count at a node is the number of occurrences of the path spelled by that node. The program checks this claim against a brute-force scan later in the lesson.

A few more facts about banana follow from the trie. The trie has 15 distinct non-empty substrings. Counting them by length gives 3 of length 1 (b, a, n), 3 of length 2 (ba, an, na), 3 of length 3 (ban, nan, ana), 3 of length 4 (bana, anan, nana), 2 of length 5 (banan, anana) and 1 of length 6 (banana). Each distinct substring has one node, and the root adds one more. So the trie for banana has 16 nodes including the root.

The program

The program below builds the suffix trie for a text. Its class SuffixTrie takes the text and the list of allowed letters, which is called the alphabet. The node stores a list with one child slot per letter, the same design as the array trie in the first lesson of this module. Each node also has pass_count, the number of suffixes that pass through it.

The class Node uses __slots__. That line tells Python to store only the two named fields in each node, with no extra dictionary for other attributes. A trie of two million nodes is much smaller this way, and the node is still used like any other object.

The first part of the main code runs the banana example. The second part builds tries for random text of four different lengths and prints the node counts and build times. The third part checks the counts against a plain scan of the text.

Save the code as suffix_demo.py, then run it with python3 suffix_demo.py.

import random
import time
 
 
class Node:
    __slots__ = ("next", "pass_count")
 
    def __init__(self, size):
        self.next = [None] * size
        self.pass_count = 0
 
 
class SuffixTrie:
    def __init__(self, text, alphabet):
        self.alphabet = alphabet
        self.root = Node(len(alphabet))
        self.node_count = 1
        for start in range(len(text)):
            self._insert_suffix(text, start)
 
    def _insert_suffix(self, text, start):
        node = self.root
        for i in range(start, len(text)):
            ch = text[i]
            idx = self.alphabet.find(ch)
            if idx < 0:
                raise ValueError(f"character not in alphabet: {ch}")
            if node.next[idx] is None:
                node.next[idx] = Node(len(self.alphabet))
                self.node_count += 1
            node = node.next[idx]
            node.pass_count += 1
 
    def count_occurrences(self, pattern):
        if pattern == "":
            raise ValueError("empty pattern")
        node = self.root
        for ch in pattern:
            idx = self.alphabet.find(ch)
            if idx < 0 or node.next[idx] is None:
                return 0
            node = node.next[idx]
        return node.pass_count
 
    def contains(self, pattern):
        return self.count_occurrences(pattern) > 0
 
 
class WholeTextOnlyTrie:
    def __init__(self):
        self.root = Node(26)
 
    def insert(self, word):
        node = self.root
        for ch in word:
            idx = ord(ch) - ord("a")
            if node.next[idx] is None:
                node.next[idx] = Node(26)
            node = node.next[idx]
 
    def contains(self, pattern):
        node = self.root
        for ch in pattern:
            idx = ord(ch) - ord("a")
            if node.next[idx] is None:
                return False
            node = node.next[idx]
        return True
 
 
def random_text(length, alphabet, seed):
    rng = random.Random(seed)
    return "".join(alphabet[rng.randrange(len(alphabet))] for _ in range(length))
 
 
alphabet = "abcdefghijklmnopqrstuvwxyz"
 
print("== Part 1: banana ==")
banana = "banana"
for i in range(len(banana)):
    print(f"suffix {i}: {banana[i:]}")
t = SuffixTrie(banana, alphabet)
print("nodes (including root):", t.node_count)
for p in ["nan", "ana", "nab", "a", "banana", "bananas"]:
    print(f"contains {p + ':':<8} {t.contains(p)}  occurrences {t.count_occurrences(p)}")
print('occurrences of "a" with uppercase input:', t.count_occurrences("A"))
 
print()
print("== Part 2: node count as the text grows (random DNA letters acgt) ==")
dna = "acgt"
for n in [250, 500, 1000, 2000]:
    text = random_text(n, dna, 7)
    start = time.perf_counter()
    big = SuffixTrie(text, dna)
    ms = int((time.perf_counter() - start) * 1000)
    print(f"n={n:5d}  nodes={big.node_count:9d}  reference slots={big.node_count * len(dna):11d}  build={ms} ms")
 
print()
print("== Part 3: a check against brute force ==")
text = random_text(300, dna, 99)
check = SuffixTrie(text, dna)
tested = 0
for length in range(1, 6):
    for start in range(0, 41 - length):
        p = text[start:start + length]
        brute = sum(1 for i in range(len(text) - length + 1) if text.startswith(p, i))
        if brute != check.count_occurrences(p):
            raise AssertionError("mismatch for " + p)
        tested += 1
print("patterns checked against a brute-force scan:", tested, "passed")

The output from one run is below. The random text comes from Python's generator with a fixed seed, so the exact letters differ from a run in another language, and the node counts differ with them. The build times depend on the machine, so yours will differ.

== Part 1: banana ==
suffix 0: banana
suffix 1: anana
suffix 2: nana
suffix 3: ana
suffix 4: na
suffix 5: a
nodes (including root): 16
contains nan:     True  occurrences 1
contains ana:     True  occurrences 2
contains nab:     False  occurrences 0
contains a:       True  occurrences 3
contains banana:  True  occurrences 1
contains bananas: False  occurrences 0
occurrences of "a" with uppercase input: 0
 
== Part 2: node count as the text grows (random DNA letters acgt) ==
n=  250  nodes=    30589  reference slots=     122356  build=12 ms
n=  500  nodes=   123449  reference slots=     493796  build=72 ms
n= 1000  nodes=   496377  reference slots=    1985508  build=328 ms
n= 2000  nodes=  1991665  reference slots=    7966660  build=1454 ms
 
== Part 3: a check against brute force ==
patterns checked against a brute-force scan: 190 passed

The contains lines show the walk at work. The pattern nab fails at the letter b, because no suffix starts with na followed by b. The pattern bananas fails at the last letter, because the trie has no path that goes past the end of banana. The pattern banana is found, because it is the first suffix. The count of ana is 2, because the suffixes anana and ana both begin with it.

The uppercase line shows a silent failure. The alphabet holds lowercase letters only, so A is not found, and the method returns 0 rather than an error. The common mistakes section explains why this is a risk.

Worked example two: occurrences and overlaps

The count of occurrences includes overlapping matches. A pattern can start again before the previous match has ended. The text mississippi shows this well. Its letters are m, i, s, s, i, s, s, i, p, p, i, numbered from position 0.

The pattern ss starts at positions 2 and 5, so the trie reports 2. The pattern issi starts at positions 1 and 4. These two matches overlap at the letter i at position 4, and the trie still counts both, which a separate run of the same class also gives as 2. The pattern i starts at positions 1, 4, 7 and 10, so the answer is 4.

A text with repeated letters shows the overlap in a sharper way. The text aaaa has the suffixes aaaa, aaa, aa and a. The trie has four distinct substrings, a, aa, aaa and aaaa, so it has 5 nodes including the root. The pattern aa starts at positions 0, 1 and 2, so the count is 3. These results come from a separate run of the same class, with this text.

Worked example three: how fast the trie grows

Part 2 of the program builds tries for random text over the four letters a, c, g and t. A DNA sequence uses these letters, which is why the alphabet is small here. The table shows the results of that run, with the node counts that include the root.

Text lengthNodesReference slotsBuild time
25030,589122,35612 ms
500123,449493,79672 ms
1,000496,3771,985,508328 ms
2,0001,991,6657,966,6601,454 ms

The table is the main result. When the text doubles, the node count grows by about four times. From 250 to 500 the count goes from 30,589 to 123,449, a factor of 4.04. From 500 to 1,000 it goes from 123,449 to 496,377, a factor of 4.02. From 1,000 to 2,000 it goes from 496,377 to 1,991,665, a factor of 4.01. That pattern means the node count grows with the square of the text length.

The reason is the suffixes. A text of length n has n suffixes, and their lengths add up to n times n + 1 divided by 2. For 2,000 letters that is 2,001,000 letter steps. Each step creates at most one node. The trie has 1,991,665 nodes, which is about 99.5 percent of the most it could have. Random letters rarely repeat a long substring, so almost every step creates a new node.

Build time also grows with the number of letter steps. Here it grew faster than the node count, because each step in Python runs several instructions. The build for 2,000 letters took about one and a half seconds on the machine used for this lesson. Memory is the real limit. The table counts reference slots, the child links held in each node's list. Each node is also a Python object, so the true memory use is higher, and a Python node costs more than one slot.

The limit is easy to estimate with arithmetic. For a text of 10,000 letters from a 26-letter alphabet, the trie can have up to 50,005,000 nodes, which is 10,000 times 10,001 divided by 2. With 26 slots per node that is about 1.3 billion reference slots, and random text of that length would come close to the upper limit. For a text of 1,000,000 letters the upper limit is about 500 billion nodes. Even at 16 bytes per node, that would take 8 terabytes.

The fix for long texts is to compress the trie, so that a chain of nodes with one child each becomes a single node. That is the idea of the radix tree, covered in lesson 09 of this module. A suffix array, which stores only the start positions in sorted order, is another common choice. Both use much less memory, and both are harder to write.

Worked example four: following one pattern

Take the pattern nana in banana. The walk has one step per letter, and each step can fail. Start at the root. It has an n child, because the suffix nana begins with n. Move to it. Its pass_count is 2, because two suffixes begin with n: nana and na. That is the number of times n occurs, which is 2.

Next, follow a. The n node has an a child, so move to it. The na node has pass_count 2 as well, because the same two suffixes begin with na. Next, follow n, then a. The nan node has pass_count 1, because only nana begins with nan. The nana node has pass_count 1, so nana occurs once, at position 2.

Now try nanab. The first four steps match, and the fifth step needs a b child under the nana node. That child does not exist, so the method returns 0 at that point. Each count can be checked by hand: n 2, na 2, nan 1, nana 1 and nanab 0.

The walk never reads the rest of banana. It reads the pattern once and stops at the first missing link. That is the reason a search is fast, even though the trie itself is large.

Common mistakes

Inserting only the whole text. A trie that holds only banana answers True for banana and False for nan, with no error. This is the most common mistake. The class WholeTextOnlyTrie in the program above is that wrong version. The output of a short program that loads it with banana is:

whole text only, contains banana: True
whole text only, contains nan: False

The fix is to insert every suffix, with a loop from each start position to the end of the text. The constructor of SuffixTrie does that.

Letters outside the alphabet in the text. The constructor raises an error for a character that the alphabet does not hold. Building the trie for Banana with a lowercase alphabet gives this message:

ValueError: character not in alphabet: B

Letters outside the alphabet in the pattern. A pattern is not checked as strictly. A pattern with a capital letter returns 0 from count_occurrences, so the result looks like "not in the text" and no error is shown. The output of the program for A shows this. Convert both the text and the pattern to one case before you build the trie or search it. In Python, text.lower() and pattern.lower() do that.

An empty pattern. The empty string has no path to walk, and the method rejects it with a message:

ValueError: empty pattern

Decide what an empty search should mean for your task before you change this check.

Reading the count as the number of distinct matches. count_occurrences counts overlapping matches, as shown with aaaa. A search that wants distinct matches needs a different rule. For example, a text that asks "how many different substrings of length 3 occur" needs the distinct count from the nodes at depth 3, not the sum of occurrences.

Estimating memory from the text length alone. A text of 1,000 letters does not need 1,000 nodes. With random DNA letters it needs about half a million. Check the table above before you decide the trie is small enough for your input.

Practice

Exercise 1. Count the occurrences of ss in mississippi, and list the start positions.

Show the answer: The count is 2. The pattern starts at position 2, where the text reads ss after the first i, and at position 5, where it reads ss after the second i.

Exercise 2. How many nodes does the suffix trie for aaaa have, including the root? List the distinct substrings it stores.

Show the answer: It has 5 nodes including the root. The distinct substrings are a, aa, aaa and aaaa, one node each, and the root adds one more. Building the trie for aaaa confirms the count of 5.

Exercise 3. Why does contains take time that depends on the length of the pattern, and not on the length of the text?

Show the answer: The trie is built before any search. A search follows one link for each letter of the pattern, from the root down, and then reads one count. The walk stops early if a link is missing. It never looks at the rest of the text, so the number of steps is at most the pattern length.

Exercise 4. Change count_occurrences so that an uppercase pattern still works on a lowercase text. Write the one line you add, and say where it goes.

Show the answer: Add pattern = pattern.lower() as the first line of the method, before the empty check. Then count_occurrences("Ana") looks for ana. The text must also be lowercase, so add text = text.lower() as the first line of the constructor. The mistakes section covers why both sides must match.

Exercise 5. A DNA text has 100,000 letters, and the alphabet has four letters. Use the formula from this lesson to estimate the largest number of nodes the suffix trie could have, and say whether it fits in memory on a typical laptop.

Show the answer: The upper limit is 100,000 times 100,001 divided by 2, which is 5,000,050,000 nodes. On a 64-bit Python each reference slot takes 8 bytes, so four slots take 32 bytes per node. That makes about 160 billion bytes, or about 160 gigabytes, before counting the Python objects themselves. Most laptops do not have that much memory, so the plain suffix trie does not fit. A radix tree or a suffix array is the usual answer at this size.

Summary

  • A suffix is a substring that runs to the end of the text, and every substring is a prefix of one suffix.
  • Inserting all suffixes turns a substring search into a prefix walk.
  • The count stored at a node gives the number of occurrences, and occurrences may overlap.
  • The node count grows with the square of the text length, which limits the plain method to short texts. Longer texts need a compressed design, such as a radix tree.
  • An uppercase pattern fails silently, while an empty pattern raises an error, so both need handling in the code.
Check yourself: Suffix tries (6 questions) →Practice this topic: coding problems for Tries →

Stuck? Ask the tutor

Ask about anything in this lesson. Answers are short and use simple words.

Sign in to ask the tutor.

← Ternary search tree versus array tries