单词拆分II

这道题考察了dp、前缀树、回溯等知识，下面给出代码

class Solution {
    class TrieNode {
        String word = null;
        HashMap<Character, TrieNode> map = new HashMap<>();
        public TrieNode() {}
    }

    public List<String> wordBreak(String s, List<String> wordDict) {
        ArrayList<String> res = new ArrayList<>();  // 保存最终结果
        if (s == null || wordDict == null) return res;
        // 构建前缀树
        TrieNode root = new TrieNode();
        for (String word : wordDict) {
            TrieNode node = root;
            for (Character letter : word.toCharArray()) {
                if (node.map.containsKey(letter)) { // 如果存在，转到下一个TrieNode
                    node = node.map.get(letter);
                } else {    // 不存在就创建一个新的TrieNode
                    TrieNode newNode = new TrieNode();
                    node.map.put(letter, newNode);
                    node = newNode;
                }
            }
            node.word = word;
        }

        int[] dp = getDp(s, root);
        backtrack(s, root, dp, 0, new ArrayList<>(), res);
        return res;
    }

    public void backtrack(String s, TrieNode root, int[] dp, int idx, ArrayList<String> words, ArrayList<String> ans) {
        if (idx == s.length()) {    // 遍历到最后一个
            ans.add(buildToSentence(words));
            return;
        }
        // 无法分解
        if (dp[idx] == 0) return;
        
        TrieNode curNode = root;
        for (int end = idx; end < s.length(); end++) {
            if (!curNode.map.containsKey(s.charAt(end))) break;
            curNode = curNode.map.get(s.charAt(end));
            if (curNode.word != null) {
                words.add(s.substring(idx, end + 1));
                backtrack(s, root, dp, end + 1, words, ans);
                words.remove(words.size() - 1); // 回溯
            }
        }
    }

    // dp[i]表示s[i...len-1]能被字典所分解的种类数
    public int[] getDp(String s, TrieNode root) {
        int len = s.length();
        int[] dp = new int[len+1];
        dp[len] = 1;    //  空串相当于1个
        
        for (int i = len-1; i >= 0; i--) {
            TrieNode node = root;
            int count = 0;
            
            for (int j = i; j < len; j++) {
                if (!node.map.containsKey(s.charAt(j))) break;
                node = node.map.get(s.charAt(j));
                if (node.word != null) count += dp[j+1];
            }
            dp[i] = count;
        }
        return dp;
    }

    // 将单词拼接成句子
    public String buildToSentence(ArrayList<String> words) {
        StringBuilder str = new StringBuilder();
        for (int idx = 0; idx < words.size(); idx++) {
            str.append(words.get(idx));
            if (idx != words.size() - 1) str.append(" ");
        }
        return str.toString();
    }
}