面试算法63：替换单词

题目

英语中有一个概念叫词根。在词根后面加上若干字符就能拼出更长的单词。例如，“an"是一个词根，在它后面加上"other"就能得到另一个单词"another”。现在给定一个由词根组成的字典和一个英语句子，如果句子中的单词在字典中有它的词根，则用它的词根替换该单词；如果单词没有词根，则保留该单词。请输出替换后的句子。
例如，如果词根字典包含字符串[“cat”，“bat”，“rat”]，英语句子为"the cattle was rattled by the battery"，则替换之后的句子是"the cat was rat by the bat"。

分析

题目中的词根其实就是前缀，因此很容易想到用前缀树来解决。用前缀树解决问题通常分为两步，第1步是创建前缀树，第2步是在前缀树中查找。

解

public class Test {
    public static void main(String[] args) {
        List<String> dict = Arrays.asList("cat", "bat", "rat");
        String sentence = "the cattle was rattled by the battery";
        String result = replaceWords(dict, sentence);
        System.out.println(result);
    }

    static class TrieNode {
        boolean isWord;
        TrieNode[] children;

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

    private static TrieNode buildTrie(List<String> dict) {
        TrieNode root = new TrieNode();
        for (String word : dict) {
            TrieNode node = root;
            for (char ch : word.toCharArray()) {
                if (node.children[ch - 'a'] == null) {
                    node.children[ch - 'a'] = new TrieNode();
                }
                node = node.children[ch - 'a'];
            }

            node.isWord = true;
        }

        return root;
    }

    private static String findPrefix(TrieNode root, String word) {
        TrieNode node = root;
        StringBuilder builder = new StringBuilder();
        for (char ch : word.toCharArray()) {
            if (node.isWord || node.children[ch - 'a'] == null) {
                break;
            }
            builder.append(ch);
            node = node.children[ch - 'a'];
        }
        return node.isWord ? builder.toString() : "";
    }

    public static String replaceWords(List<String> dict, String sentence) {
        TrieNode root = buildTrie(dict);
        StringBuilder builder = new StringBuilder();

        String[] words = sentence.split(" ");
        for (int i = 0; i < words.length; i++) {
            String prefix = findPrefix(root, words[i]);
            if (!prefix.isEmpty()) {
                words[i] = prefix;
            }
        }

        return String.join(" ", words);
    }
}