Trie Implementation in JAVA

This is my Java implementation of the Trie data structure using the Java sorted map data structure, TreeMap, in the TrieNode class. I am using the TreeMap DS instead of the HashMap because I would like to maintain the reverse alphabetical order of the entries. The methods the Trie class supports are add, getWordCount, getWords, and unsetWord. The getWords method is iteratively implemented using depth first traversal.

public class TrieNode {
    char data;
    Map<Character, TrieNode> children;
    boolean isWord;
    int wordCount;

    public TrieNode(char k) {
        data = k;
        children = new TreeMap<Character, TrieNode>(Collections.reverseOrder());
        isWord = false;
        wordCount = 0;
    }
}

public class Trie {
    TrieNode root;

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

    public void add(String word) {
        int n = word.length();

        if (n == 0) return;

        TrieNode cur = root;
        for (int i = 0; i < n; i++) {
            char c = word.charAt(i);

            if (cur.children.get(c) == null)
                cur.children.put(c, new TrieNode(c));

            cur = cur.children.get(c);
            cur.wordCount++;
        }

        cur.isWord = true;
    }

    public int getWordCount(String prefix) {
        int n = prefix.length();

        if (n == 0) return 0;

        TrieNode cur = root;
        for (int i = 0; cur != null && i < n; i++) {
            char c = prefix.charAt(i);
            cur = cur.children.get(i);
        }

        return cur.wordCount;
    }

    public List<String> getWords(String prefix) {
        int n = prefix.length();

        if (n == 0) return new ArrayList<String>();

        TrieNode cur = root;
        for (int i = 0; cur != null && i < n; i++) {
            char c = prefix.charAt(i);
            cur = cur.children.get(c);
        }

        if (cur == null || cur.wordCount == 0) return new ArrayList<String>();

        List<String> words = new ArrayList<String>();
        Stack<TrieNode> parents = new Stack<TrieNode>();
        Stack<Integer> levels = new Stack<Integer>();
        StringBuilder sb = new StringBuilder();
        int level = prefix.length() - 1;

        parents.push(cur);
        levels.push(level);

        while (!parents.isEmpty()) {
            cur = parents.pop();
            level = levels.pop();

            sb.setLength(level);
            sb.append(cur.data);

            if (cur.isWord) words.add(sb.toString());

            Iterator<Character> children = cur.children.keySet().iterator();

            while (children.hasNext()) {
                char c = children.next();
                TrieNode child = cur.children.get(c);
                parents.push(child);
                levels.push(level + 1);
            }
        }

        return words;
    }

    public void unsetWord(String word) {
        int n = word.length();
        if (n == 0) return;

        TrieNode cur = root;
        for (int i = 0; cur != null && i < n; i++) {
            char c = word.charAt(i);
            cur = cur.children.get(c);
        }

        if (cur == null) return;

        cur.isWord = false;
    }
}

This article is my 8th oldest. It is 375 words long, and it’s got 0 comments for now.