Total Number of Words in a Trie

Given a trie, return the number of distinct words stored in it.

Example

Text

insert: "the", "a", "there", "answer", "any", "by", "bye", "their", "the"
Output: 8     ("the" was inserted twice, but it's one word)

graph TB
    root(("root")) -->|t| t["t"]
    root -->|a| a["a (end)"]
    root -->|b| b["b"]
    t -->|h| th["th"]
    th -->|e| the["the (end)"]
    the -->|r| ther["ther"]
    ther -->|e| there["there (end)"]
    the -->|i| thei["thei"]
    thei -->|r| their["their (end)"]
    a -->|n| an["an"]
    an -->|"s, w, e, r"| answer["answer (end)"]
    an -->|y| any["any (end)"]
    b -->|y| by["by (end)"]
    by -->|e| bye["bye (end)"]

Every word ends at exactly one node, and that node has isEnd set. So the number of words is the number of nodes with isEnd == true. There are 8 of them.

Brute force: count end-of-word nodes

Traverse the whole trie with a DFS and count the nodes whose isEnd flag is set. A trie node is an object with an array of 26 child references. A slot is null until a word needs that child.

Java

public class Main {
    static class TrieNode {
        TrieNode[] children = new TrieNode[26];
        boolean isEnd;
    }

    static void insert(TrieNode root, String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            int idx = word.charAt(i) - 'a';
            if (node.children[idx] == null) node.children[idx] = new TrieNode();
            node = node.children[idx];
        }
        node.isEnd = true;
    }

    static int countWords(TrieNode node) {
        int total = node.isEnd ? 1 : 0; // 1 if a word ends here, else 0
        for (TrieNode child : node.children) {
            if (child != null) total += countWords(child);
        }
        return total;
    }

    public static void main(String[] args) {
        TrieNode root = new TrieNode();
        for (String w : new String[] {"the", "a", "there", "answer", "any", "by", "bye", "their", "the"}) {
            insert(root, w);
        }
        System.out.println(countWords(root));
    }
}

Output

Time: O(N · 26), where N is the number of nodes in the trie, since every node checks its 26 slots. Space: O(L) for the recursion, where L is the length of the longest word.

Optimal: count while inserting

Keep a words counter in the Trie. When insert reaches the last node of a word, increment the counter only if that node's isEnd flag wasn't already set. Otherwise, the word is a duplicate.

Walkthrough of the last few inserts:

insert	`isEnd` before	counter after
... 7 distinct words ...	false each time	7
"their"	false	8
"the" (again)	true	8 (unchanged)

Java

public class Main {
    static class TrieNode {
        TrieNode[] children = new TrieNode[26];
        boolean isEnd;
    }

    static class Trie {
        TrieNode root = new TrieNode();
        int words;

        void insert(String word) {
            TrieNode node = root;
            for (int i = 0; i < word.length(); i++) {
                int idx = word.charAt(i) - 'a';
                if (node.children[idx] == null) node.children[idx] = new TrieNode();
                node = node.children[idx];
            }
            if (!node.isEnd) {
                node.isEnd = true;
                words++; // a new word, not a duplicate
            }
        }

        int size() {
            return words;
        }
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        for (String w : new String[] {"the", "a", "there", "answer", "any", "by", "bye", "their", "the"}) {
            trie.insert(w);
        }
        System.out.println(trie.size());
    }
}

Output

Time: O(1) to answer. insert stays O(L). Space: O(1) extra: one counter.

The check if (!node.isEnd) before setting the flag is the "was this word already here?" test. If the trie also supports deletion, decrement the counter whenever a removal succeeds.

Edge cases

Empty trie: 0.
Duplicate inserts: counted once.
A word that's a prefix of another (the and there): both are counted, since each has its own end node.