Trie/ Leetcode 648. Replace Word

Relevant problem: Leetcode 648. Replace Word
Using Trie, we can search the key in $O(M)$ time, with a penalty on Trie storage requirement.

An implementation:

class TrieNode {
    // Initialize your data structure here.
    public TrieNode[] children;
    public char curChar;
    public boolean isLeaf;

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

public class Trie {
    private TrieNode root;

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

    // Inserts a word into the trie.
    public void insert(String word) {
        int len = word.length();
        if(len == 0) {
            return;
        }
        TrieNode curNode = root;
        for(int i = 0; i < len; i++) {
            char cur = word.charAt(i);
            int index = cur - 'a';
            //if there is no such a node, add a new one
            if(curNode.children[index] == null) {
                curNode.children[index] = new TrieNode();
                curNode.children[index].curChar = cur;
            }
            curNode = curNode.children[index];
        }
        curNode.isLeaf = true;
    }

    // Returns if the word is in the trie.
    public boolean search(String word) {
        if(word.length() == 0) {
            return true;
        }
        TrieNode curNode = root;
        for(int i = 0; i < word.length(); i++) {
            char cur = word.charAt(i);
            int index = cur - 'a';
            if(curNode.children[index] == null) {
                return false;
            }
            curNode = curNode.children[index];
        }
        return curNode.isLeaf;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    public boolean startsWith(String prefix) {
        if(prefix.length() == 0) {
            return true;
        }
        TrieNode curNode = root;
        for(int i = 0; i < prefix.length(); i++) {
            char cur = prefix.charAt(i);
            int index = cur - 'a';
            if(curNode.children[index] == null) {
                return false;
            }
            if(curNode.children[index].curChar != cur) {
                return false;
            }
            curNode = curNode.children[index];
        }
        return true;
    }
}

And here is a sample solution for the Leetcode Problem

public String replaceWords(List<String> dict, String sentence) {
        String[] tokens = sentence.split(" ");
        TrieNode trie = buildTrie(dict);
        return replaceWords(tokens, trie);
    }

    private String replaceWords(String[] tokens, TrieNode root) {
        StringBuilder stringBuilder = new StringBuilder();
        for (String token : tokens) {
            stringBuilder.append(getShortestReplacement(token, root));
            stringBuilder.append(" ");
        }
        return stringBuilder.substring(0, stringBuilder.length()-1);
    }

    private String getShortestReplacement(String token, final TrieNode root) {
        TrieNode temp = root;
        StringBuilder stringBuilder = new StringBuilder();
        for (char c : token.toCharArray()) {
            stringBuilder.append(c);
            if (temp.children[c - 'a'] != null) {
                if (temp.children[c - 'a'].isWord) {
                    return stringBuilder.toString();
                }
                temp = temp.children[c - 'a'];
            } else {
                return token;
            }
        }
        return token;
    }

    private TrieNode buildTrie(List<String> dict) {
        TrieNode root = new TrieNode(' ');
        for (String word : dict) {
            TrieNode temp = root;
            for (char c : word.toCharArray()) {
                if (temp.children[c - 'a'] == null) {
                    temp.children[c - 'a'] = new TrieNode(c);
                }
                temp = temp.children[c - 'a'];
            }
            temp.isWord = true;
        }
        return root;
    }

    public class TrieNode {
        char val;
        TrieNode[] children;
        boolean isWord;

        public TrieNode(char val) {
            this.val = val;
            this.children = new TrieNode[26];
            this.isWord = false;
        }
    }