Implement Trie (Prefix Tree) medium

Problem Statement

Implement a Trie (Prefix Tree) data structure.

A trie is a tree-like data structure used for storing strings efficiently, allowing for fast retrieval of strings and prefixes of strings. Each node in a trie represents a character in a string, and paths from the root to a node represent prefixes of strings stored in the trie. Implement the Trie class:

  • Trie(): Initializes the trie object.
  • void insert(String word): Inserts the string word into the trie.
  • boolean search(String word): Returns true if the string word is in the trie (i.e., was inserted before), and false otherwise.
  • boolean startsWith(String prefix): Returns true if there is any string in the trie that starts with the given prefix, and false otherwise.

Example 1

Input
["Trie", "insert", "search", "search", "startsWith", "insert", "search"]
[[], ["apple"], ["apple"], ["app"], ["app"], ["app"], ["app"]]
Output
[null, null, true, false, true, null, true]

Explanation
Trie trie = new Trie();
trie.insert("apple");
trie.search("apple");   // return True
trie.search("app");     // return False
trie.startsWith("app"); // return True
trie.insert("app");
trie.search("app");     // return True

Example 2

Input
["Trie", "insert", "insert", "search", "search", "startsWith", "startsWith"]
[[], ["a"], ["app"], ["app"], ["ap"], ["ap"], ["a"]]
Output
[null,null,null,true,false,true,true]

Explanation
Trie trie = new Trie();
trie.insert("a");
trie.insert("app");
trie.search("app");     // return True
trie.search("ap");      // return False
trie.startsWith("ap");  // return True
trie.startsWith("a");   // return True

Steps

  1. Trie Node Class: Create a TrieNode class to represent each node in the trie. Each node should have:

    • An array of TrieNodes (children) representing characters 'a' to 'z'.
    • A boolean variable isEnd to indicate whether this node represents the end of a word.
  2. Trie Class: Create a Trie class with a root TrieNode.

  3. insert(String word):

    • Iterate through the characters of the word.
    • For each character, if there is no child node for that character, create one.
    • Move to the child node.
    • After processing all characters, set isEnd to true for the last node.
  4. search(String word):

    • Iterate through the characters of the word.
    • If at any point a child node for the current character is not found, return false.
    • After processing all characters, return true only if the last node has isEnd set to true.
  5. startsWith(String prefix):

    • Similar to search, but you don't need to check isEnd at the end. Return true if you reach the end of the prefix successfully, indicating that there's at least one word starting with that prefix.

Explanation

The Trie data structure excels at prefix-based operations because it organizes words based on their prefixes. Each branch of the Trie represents a possible prefix. Searching and checking for prefixes becomes a matter of traversing down the tree. The isEnd flag efficiently determines if a complete word has been found.

Code

class TrieNode {
    TrieNode[] children;
    boolean isEnd;

    public TrieNode() {
        children = new TrieNode[26];
        isEnd = false;
    }
}

class Trie {
    TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    public void insert(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                node.children[index] = new TrieNode();
            }
            node = node.children[index];
        }
        node.isEnd = true;
    }

    public boolean search(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                return false;
            }
            node = node.children[index];
        }
        return node.isEnd;
    }

    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (char c : prefix.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                return false;
            }
            node = node.children[index];
        }
        return true;
    }
}

Complexity

  • Time Complexity:

    • insert: O(m), where m is the length of the word.
    • search: O(m), where m is the length of the word.
    • startsWith: O(m), where m is the length of the prefix.
  • Space Complexity: O(N*M), where N is the number of words inserted, and M is the maximum length of a word. (This is because in the worst case, you might have a separate branch for each character in each word.)