Trie Data Structure Cheat Sheet in Java

1615

Trie Data Structure Cheat Sheet in Java

A Trie (also known as a Prefix Tree) is a tree-based data structure that is used to store a dynamic set of strings, where the keys are usually strings. It is particularly useful for tasks such as:

Searching for words with a common prefix
Autocomplete
Spell checking
Dictionary-based searching

Basic Structure of a Trie

A Trie typically has the following components:

Nodes: Each node in the Trie represents a single character of a word.
Children: Each node contains a map or array to store child nodes.
End of Word: A boolean value to mark if the node represents the end of a word.

Trie Node Class

class TrieNode {
    // HashMap to store children nodes
    Map<Character, TrieNode> children;
    boolean isEndOfWord;

    public TrieNode() {
        children = new HashMap<>();
        isEndOfWord = false;
    }
}

Trie Class with Basic Operations

1. Insert Operation

Inserting a word into the Trie involves adding characters one by one. If the character already exists as a child node, move to that node; otherwise, create a new node for that character.

class Trie {
    private TrieNode root;

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

    // Insert a word into the trie
    public void insert(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            node = node.children.computeIfAbsent(c, k -> new TrieNode());
        }
        node.isEndOfWord = true;
    }
}

2. Search Operation

Searching for a word in the Trie involves traversing the Trie node-by-node following the word's characters. If the traversal reaches the end of the word and the node is marked as isEndOfWord = true, then the word exists in the Trie.

// Search if the word exists in the trie
public boolean search(String word) {
    TrieNode node = root;
    for (char c : word.toCharArray()) {
        node = node.children.get(c);
        if (node == null) return false;
    }
    return node.isEndOfWord;
}

3. Starts With Prefix Operation

This operation checks whether any word in the Trie starts with a given prefix. It only needs to traverse the characters of the prefix, not the whole word.

// Check if there's any word in the trie that starts with the given prefix
public boolean startsWith(String prefix) {
    TrieNode node = root;
    for (char c : prefix.toCharArray()) {
        node = node.children.get(c);
        if (node == null) return false;
    }
    return true;
}

4. Delete Operation (Optional)

To delete a word, you need to remove nodes that are part of the word and ensure that the nodes are removed only if they are not part of any other word. This is a bit tricky and requires backtracking.

Trie Example

Let’s put everything together into a full example.

import java.util.HashMap;
import java.util.Map;

class TrieNode {
    Map<Character, TrieNode> children;
    boolean isEndOfWord;

    public TrieNode() {
        children = new HashMap<>();
        isEndOfWord = false;
    }
}

class Trie {
    private TrieNode root;

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

    // Insert a word into the trie
    public void insert(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            node = node.children.computeIfAbsent(c, k -> new TrieNode());
        }
        node.isEndOfWord = true;
    }

    // Search if the word exists in the trie
    public boolean search(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            node = node.children.get(c);
            if (node == null) return false;
        }
        return node.isEndOfWord;
    }

    // Check if there's any word in the trie that starts with the given prefix
    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (char c : prefix.toCharArray()) {
            node = node.children.get(c);
            if (node == null) return false;
        }
        return true;
    }

    // Delete a word from the trie
    public boolean delete(String word) {
        return deleteHelper(root, word, 0);
    }

    private boolean deleteHelper(TrieNode node, String word, int index) {
        if (index == word.length()) {
            if (!node.isEndOfWord) return false;  // Word not found
            node.isEndOfWord = false;
            return node.children.isEmpty();  // If no children, it can be deleted
        }

        char ch = word.charAt(index);
        TrieNode childNode = node.children.get(ch);
        if (childNode == null) return false;  // Word not found

        boolean shouldDeleteChild = deleteHelper(childNode, word, index + 1);

        if (shouldDeleteChild) {
            node.children.remove(ch);
            return node.children.isEmpty() && !node.isEndOfWord;  // If the node can be deleted
        }
        return false;
    }
}

public class Main {
    public static void main(String[] args) {
        Trie trie = new Trie();

        // Insert words
        trie.insert("apple");
        trie.insert("app");
        trie.insert("bat");
        trie.insert("batman");

        // Search words
        System.out.println(trie.search("apple"));  // true
        System.out.println(trie.search("app"));    // true
        System.out.println(trie.search("bat"));    // true
        System.out.println(trie.search("bats"));   // false

        // Prefix search
        System.out.println(trie.startsWith("ba")); // true
        System.out.println(trie.startsWith("cat")); // false

        // Delete words
        trie.delete("batman");
        System.out.println(trie.search("batman")); // false
        System.out.println(trie.search("bat"));    // true
    }
}

Operations Summary

Operation	Time Complexity	Space Complexity
`insert(word)`	O(L)	O(L)
`search(word)`	O(L)	O(L)
`startsWith(prefix)`	O(L)	O(L)
`delete(word)`	O(L)	O(L)

Where:

L is the length of the word (or prefix).
Space Complexity is mainly driven by the size of the Trie, which is proportional to the number of characters across all inserted words.

Common Use Cases for Tries

Autocomplete: Tries are great for autocomplete functionality, where you need to find all words starting with a given prefix.
Spell-checking: Checking whether a word is valid based on a dictionary.
Prefix Matching: Efficiently finding all words starting with a particular prefix.
Pattern Matching: Tries can be extended for solving string matching problems (e.g., wildcard matching).

Summary

A Trie is an efficient tree-like structure for storing strings.
It supports operations like insertion, search, and prefix search in O(L) time, where L is the length of the word.
It uses extra space, but it's highly efficient for operations that involve a lot of string searches, especially when dealing with prefixes.

This cheat sheet provides the basic structure, operations, and a full example to help you get started with implementing Tries in Java!

Comments (1)