[leetCode]425. 单词方块

题目

https://leetcode-cn.com/problems/word-squares/

字典树 + 回溯

class Solution {

    private List<List<String>> ans = new ArrayList<>();
    private List<String> ansBuilder = new ArrayList<>();

    public List<List<String>> wordSquares(String[] words) {
        if (words == null || words.length == 0)
            return ans;
        
        Trie trie = new Trie(words);
        
        for (String w : words) {
            int len = w.length();
            ansBuilder.add(w);
            backTracking(len, trie);
            ansBuilder.remove(ansBuilder.size() - 1);
        }
        return ans;
    }

    private void  backTracking(int len, Trie trie) {
        if (ansBuilder.size() == len) {
            ans.add(new ArrayList<>(ansBuilder));
            return;
        }
        // 前缀的长度
        int i = ansBuilder.size();
        StringBuilder prefix = new StringBuilder();
        for (String s : ansBuilder) {
            prefix.append(s.charAt(i));
        }
        // 符合该前缀的字符串
        List<String> startWith = trie.findByPrefix(prefix.toString());
        for (String s : startWith) {
            ansBuilder.add(s);
            backTracking(len, trie);
            ansBuilder.remove(ansBuilder.size() - 1);
        }
    }
    
    class TrieNode {
        List<String> startWith;
        Map<Character, TrieNode> children;

        public TrieNode() {
            startWith = new ArrayList<>();
            children = new HashMap<>();
        }
    }

    class Trie {
        TrieNode root;
         
        public Trie(String[] words) {
            root = new TrieNode();
            for (String w : words) {
                TrieNode cur = root;
                for (char ch : w.toCharArray()) {
                    TrieNode next = cur.children.computeIfAbsent(ch, k -> new TrieNode());
                    next.startWith.add(w);
                    cur = next;
                }
            }
        }

        public List<String> findByPrefix(String prefix) {
            List<String> ans = new ArrayList<>();
            TrieNode cur = root;
            for (char ch : prefix.toCharArray()) {
                if (!cur.children.containsKey(ch)) {
                    return ans;
                }
                cur = cur.children.get(ch);
            }
            ans.addAll(cur.startWith);
            return ans;
        }
    }
}

数组实现字典树

class Solution {
    List<List<String>> ans = new ArrayList<>();
    List<String> ansbuilder = new ArrayList<>();
    TrieNode root;
    public List<List<String>> wordSquares(String[] words) {
        root = new TrieNode();
        for (String word : words) {
            insert(word);
        }
       
        for (String word : words) {
            ansbuilder.add(word);
            backTracking(word.length(), root);
            ansbuilder.remove(ansbuilder.size() - 1);
        }
        return ans;
    }

    void backTracking(int len, TrieNode node) {
        if (ansbuilder.size() == len) {
            ans.add(new ArrayList<>(ansbuilder));
            return;
        }
        int idx = ansbuilder.size();
        StringBuilder prefix = new StringBuilder();
        for (String word : ansbuilder) {
            prefix.append(word.charAt(idx));
        }
        List<String> words = findStartWith(prefix.toString());
        for (String word : words) {
            ansbuilder.add(word);
            backTracking(len, node);
            ansbuilder.remove(ansbuilder.size() - 1);
        }
    }

    class TrieNode {
        List<String> words;
        TrieNode[] children;
        TrieNode() {
            words = new ArrayList<>();
            children = new TrieNode[26];
        }
    }

    public void insert(String word) {
        TrieNode cur = root;
        for (char c : word.toCharArray()) {
            if (cur.children[c - 'a'] == null)
                cur.children[c - 'a'] = new TrieNode();
            cur = cur.children[c - 'a'];
            cur.words.add(word);
        }
    }
    public List<String> findStartWith(String word) {
        List<String> ans = new ArrayList<>();
        TrieNode cur = root;
        for (char c : word.toCharArray()) {
            if (cur.children[c - 'a'] == null) {
                return ans;
            }
            cur = cur.children[c - 'a'];
        }
        ans.addAll(cur.words);
        return ans;
    }
}