Implement Trie (Prefix Tree) - Java Solution

1. Introduction

In this post, we'll explore how to implement a Trie, also known as a Prefix Tree. A Trie is a specialized tree-like data structure that is particularly useful for managing a dynamic set of strings where keys are usually strings. Common applications include autocomplete systems and spell checkers.

Problem

The task is to implement the Trie class with the following functions:

- Trie(): Initializes the trie object.

- void insert(String word): Inserts the string word into the trie.

- boolean search(String word): Returns true if word is in the trie.

- boolean startsWith(String prefix): Returns true if there is a word in the trie that starts with the given prefix.

2. Solution Steps

1. Define the TrieNode class representing each node of the Trie.

2. Each TrieNode contains an array of children TrieNodes and a boolean flag indicating the end of a word.

3. Implement the Trie constructor to initialize the root of the Trie.

4. insert: For each character in the word, travel down the Trie, creating new nodes as needed. Mark the end of the word in the Trie.

5. search: Traverse the Trie with the given word and check if the word exists, considering the end of the word flag.

6. startsWith: Similar to search, but only check if the prefix exists in the Trie, without considering the end of the word flag.

3. Code Program

class TrieNode {
    private TrieNode[] children;
    private boolean isEndOfWord;

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

    public boolean containsKey(char ch) {
        return children[ch - 'a'] != null;
    }

    public TrieNode get(char ch) {
        return children[ch - 'a'];
    }

    public void put(char ch, TrieNode node) {
        children[ch - 'a'] = node;
    }

    public void setEnd() {
        isEndOfWord = true;
    }

    public boolean isEnd() {
        return isEndOfWord;
    }
}

class Trie {
    private TrieNode root;

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

    public void insert(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            if (!node.containsKey(ch)) {
                node.put(ch, new TrieNode());
            }
            node = node.get(ch);
        }
        node.setEnd();
    }

    public boolean search(String word) {
        TrieNode node = searchPrefix(word);
        return node != null && node.isEnd();
    }

    public boolean startsWith(String prefix) {
        TrieNode node = searchPrefix(prefix);
        return node != null;
    }

    private TrieNode searchPrefix(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            if (node.containsKey(ch)) {
                node = node.get(ch);
            } else {
                return null;
            }
        }
        return node;
    }
}

public class Main {
    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("apple");
        System.out.println(trie.search("apple"));   // true
        System.out.println(trie.search("app"));     // false
        System.out.println(trie.startsWith("app")); // true
        trie.insert("app");
        System.out.println(trie.search("app"));     // true
    }
}

Output:

true
false
true
true

Explanation:

1. TrieNode class: Represents each node in the Trie. Each node contains an array children for 26 lowercase English letters and a boolean isEndOfWord.

2. Trie class: Implements the Trie with functions insert, search, and startsWith.

3. insert: Iteratively creates nodes for each character in the word and marks the end of the word.

4. search: Searches for the word in the Trie and checks if it ends at the node (isEndOfWord).

5. startsWith: Similar to search, but only checks if the prefix exists.

6. searchPrefix: A helper function used by both search and startsWith to traverse the Trie according to the given word or prefix.

7. The implementation ensures that each operation of the Trie, namely insert, search, and startsWith, is done in O(n) time complexity, where n is the key length.