Number of Distinct Substrings in a String | Trie | C++ | Java

Finding the total number of distinct substrings in a given string is a classic algorithmic challenge. While brute-force methods scale quadratically or cubically, leveraging a Trie (Prefix Tree) data structure allows us to compute the exact count of unique substrings with optimal time efficiency.

1. Understanding the Problem: Distinct Substrings

Given a string (e.g., abab), our goal is to find all unique substrings that can be formed. A substring is any contiguous sequence of characters within a string.

For example, consider the string abab:

• Substrings of length 1: a, b, a, b
• Substrings of length 2: ab, ba, ab
• Substrings of length 3: aba, bab
• Substrings of length 4: abab

Listing every single substring naively and storing them in a hash set results in heavy memory and time consumption. We need a systematic way to generate and count them without redundant checks.

2. Naive Approach Using Hash Sets

The most intuitive brute-force method is to generate all possible substrings using nested loops and insert them into a HashSet. Because a set automatically filters out duplicates, the size of the set at the end will give the number of distinct substrings.

Complexity Analysis of the Naive Approach

• Time Complexity: $\mathcal{O}(N^2 \times \text{length})$ or $\mathcal{O}(N^3)$, where $N$ is the length of the string. Generating all substrings takes $\mathcal{O}(N^2)$ loops, and inserting strings of length up to $N$ into a hash set adds an extra factor.
• Space Complexity: $\mathcal{O}(N^3)$ in the worst case because we store up to $\frac{N(N+1)}{2}$ substrings, each taking up to $\mathcal{O}(N)$ space in memory.

3. Core Mental Model: The Trie Approach

To optimize both time and space, we can visualize all substrings of a string as prefixes of its suffixes.

If you take any string and list all its suffixes:

• String: abab
• Suffix 1: abab
• Suffix 2: bab
• Suffix 3: ab
• Suffix 4: b

Every single substring of the original string is simply a prefix of one of these suffixes. Therefore, the problem of counting distinct substrings transforms into:

"Insert all suffixes of the string into a Trie. Every time a new node is created during insertion, it represents a unique prefix—and thus a unique substring."

4. Step-by-Step Flow Walkthrough

Let's trace how we insert suffixes into a Trie for the string abab.

1. Generate Suffixes:

   • Start index 0: abab
   • Start index 1: bab
   • Start index 2: ab
   • Start index 3: b

2. Insert Suffixes into Trie:

   • Insert abab: creates nodes for a, b, a, b (4 new nodes).
   • Insert bab: b already exists, but creates new nodes for a, b (2 new nodes).
   • Insert ab: a, b both exist (0 new nodes).
   • Insert b: b already exists (0 new nodes).

Total new nodes created = 4 + 3 + 0 + 0 (plus root offset or tracking counter) = exact count of distinct substrings.

5. Trie Node Structure Design

To implement this in C++ or Java, our Trie node needs a reference to its children (representing characters a through z) and a counter or boolean flag.

class Node {
    Node* links[26];
    
    bool containsKey(char ch) {
        return links[ch - 'a'] != nullptr;
    }
    
    void put(char ch, Node* node) {
        links[ch - 'a'] = node;
    }
    
    Node* get(char ch) {
        return links[ch - 'a'];
    }
};

6. Complete Algorithm Implementation

Here is how the complete algorithm operates to count distinct substrings by iterating through every starting index, generating suffixes on the fly, and inserting them into the Trie.

#include <iostream>
#include <string>

using namespace std;

struct Node {
    Node* links[26];

    bool containsKey(char ch) {
        return links[ch - 'a'] != nullptr;
    }

    void put(char ch, Node* node) {
        links[ch - 'a'] = node;
    }

    Node* get(char ch) {
        return links[ch - 'a'] = node; // wait, standard getter: return links[ch - 'a'];
    }
};

class Solution {
public:
    int countDistinctSubstrings(string s) {
        Node* root = new Node();
        int distinctCount = 0;
        int n = s.length();

        for (int i = 0; i < n; i++) {
            Node* node = root;
            for (int j = i; j < n; j++) {
                // If character does not exist, create a new node and increment count
                if (!node->containsKey(s[j])) {
                    node->put(s[j], new Node());
                    distinctCount++;
                }
                node = node->get(s[j]);
            }
        }
        return distinctCount + 1; // +1 accounts for the empty substring
    }
};

Note: Depending on whether your problem definition includes the empty string "" as a distinct substring, you may add +1 to the final count.

7. Complexity Analysis: Trie vs. Naive

Summary and Key Takeaways

• Problem Definition: Finding distinct substrings requires filtering out duplicate sequences efficiently across all lengths ($1$ to $N$).
• The Suffix Connection: Every substring of a string is a prefix of at least one of its suffixes.
• Trie Advantage: Inserting suffixes into a Trie naturally groups shared prefixes together.
• Node Creation Metric: Each time a new Trie node is instantiated during suffix insertion, it corresponds directly to a newly discovered distinct substring.
• Time Efficiency: Reduces time complexity from cubic $\mathcal{O}(N^3)$ down to quadratic $\mathcal{O}(N^2)$.
• Space Optimization: Eliminates the need to allocate memory for full substring copies in a HashSet, utilizing a compact tree of character pointers instead.