Java による Trie 木の実装

leetCode の「Implement Trie (Prefix Tree)」を解いた。
https://leetcode.com/problems/implement-trie-prefix-tree

今後、Trie 木を実装する際のヒントにしたいので、作成したプログラムをここに残しておく。

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

class Trie {

    TrieNode root = new TrieNode();

    /** Initialize your data structure here. */
    public Trie() {

    }

    /** Inserts a word into the trie. */
    public void insert(String word) {
        if (word == null || word.length() == 0) {
            return;
        }
        TrieNode currNode = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            TrieNode child = currNode.getChildNodes().get(ch);
            if (child == null) {
                TrieNode newNode = new TrieNode();
                currNode.getChildNodes().put(ch, newNode);
                currNode = newNode;
            }
            else {
                currNode = child;
            }
        }
        currNode.setWord(true);
    }

    /** Returns if the word is in the trie. */
    public boolean search(String word) {
        if (word == null || word.length() == 0) {
            return false;
        }
        TrieNode node = searchPrefix(word, root);
        if (node == null) {
            return false;
        }
        return node.isWord();
    }

    /** Returns if there is any word in the trie that starts with the given prefix. */
    public boolean startsWith(String prefix) {
        if (prefix == null || prefix.length() == 0) {
            return false;
        }
        TrieNode node = searchPrefix(prefix, root);
        if (node == null) {
            return false;
        }
        return true;
    }

    private TrieNode searchPrefix(String prefix, TrieNode root) {
        if (prefix == null || prefix.length() == 0 || root == null) {
            return null;
        }
        TrieNode currNode = root;
        for (char ch : prefix.toCharArray()) {
            TrieNode child = currNode.getChildNodes().get(ch);
            if (child == null) {
                return null;
            }
            currNode = child;
        }
        return currNode;
    }

    private static class TrieNode {
        //「後続の文字 → 後続の文字に対応する子ノード」のマップ
        // 例えば、「root -> "a" -> "p" -> "p" -> "l" -> "e"」というツリーがあったとして、
        // "a" のノードにいる時、"p" を key とする子ノードがあれば、"ap" までは持っていると判断する。
        private Map<Character, TrieNode> childNodes = new HashMap<>();

        // このノードが単語(終端)かどうかのフラグ
        private boolean isWord = false;

        private Map<Character, TrieNode> getChildNodes() {
            return childNodes;
        }

        private boolean isWord() {
            return isWord;
        }

        private void setWord(boolean word) {
            isWord = word;
        }
    }
}

ちなみに、上記を leetCode に提出すると、以下の結果となった。

Runtime: 37 ms, faster than 43.03% of Java online submissions for Implement Trie (Prefix Tree).
Memory Usage: 51.4 MB, less than 96.15% of Java online submissions for Implement Trie (Prefix Tree).