【Lintcode】634. Word Squares

题目地址：

https://www.lintcode.com/problem/word-squares/description

对于一个若干单词组成的数组，对于任意的$i$，如果第$i$行和第$i$列念出来的单词是一样的，则称其为一个单词方阵。给定若干单词组成的数组$A$，求所有$A$中的单词组成的单词方阵。题目保证单词只含英文小写字母，长度都相同，并且没有重复单词。

思路是DFS + 剪枝。其实就是直接枚举每个位置可以填哪些单词即可，没有别的好的办法。但我们仍然可以做若干剪枝操作。首先，

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Solution {
    
    // 开一个前缀树的类，可以实现插入单词，和得到某个字符串为前缀的所有单词，这两个功能
    class Trie {
        class Node {
            boolean isWord;
            Node[] nexts;
            
            public Node() {
                nexts = new Node[26];
            }
        }
        
        private Node root;
        // key是字符串，value是以key为前缀的所有单词组成的list。该map是做缓存之用
        private Map<String, List<String>> map;
        
        public Trie() {
            root = new Node();
            map = new HashMap<>();
        }
        
        public Trie(String[] words) {
            this();
            for (String word : words) {
                insert(word);
            }
        }
        
        public void insert(String s) {
            Node cur = root;
            for (int i = 0; i < s.length(); i++) {
                char ch = s.charAt(i);
                if (cur.nexts[ch - 'a'] == null) {
                    cur.nexts[ch - 'a'] = new Node();
                }
                
                cur = cur.nexts[ch - 'a'];
            }
            
            cur.isWord = true;
        }
        
        // 得到以pref为前缀的所有单词
        public List<String> getPrefix(String pref) {
        	// 如果缓存里有该前缀，则直接拿出来
            if (map.containsKey(pref)) {
                return map.get(pref);
            }
            
            List<String> res = new ArrayList<>();
            StringBuilder sb = new StringBuilder();
            Node cur = root;
            for (int i = 0; i < pref.length(); i++) {
                char ch = pref.charAt(i);
                // 如果不存在该前缀的单词，则直接返回空列表
                if (cur.nexts[ch - 'a'] == null) {
                    return res;
                }
                
                sb.append(ch);
                cur = cur.nexts[ch - 'a'];
            }
            
            dfs(cur, sb, res);
            map.putIfAbsent(pref, res);
            
            return res;
        }
        
        private void dfs(Node cur, StringBuilder sb, List<String> res) {
            if (cur.isWord) {
                res.add(sb.toString());
            }
            
            for (char ch = 'a'; ch <= 'z'; ch++) {
                if (cur.nexts[ch - 'a'] != null) {
                    sb.append(ch);
                    dfs(cur.nexts[ch - 'a'], sb, res);
                    sb.setLength(sb.length() - 1);
                }
            }
        }
    }
    
    /*
     * @param words: a set of words without duplicates
     * @return: all word squares
     */
    public List<List<String>> wordSquares(String[] words) {
        // write your code here
        if (words == null || words.length == 0) {
            return res;
        }
        
        Trie trie = new Trie(words);
        List<List<String>> res = new ArrayList<>();

        dfs(0, new ArrayList<>(), words, res, trie);
        
        return res;
    }
    
    // pos记录枚举到的列的下标，cur存已经枚举的单词
    private void dfs(int pos, List<String> cur, String[] words, List<List<String>> res, Trie trie) {
    	// 如果所有列都已经枚举完了，说明找到了一个合法解，将其加入res
        if (pos == words[0].length()) {
            res.add(new ArrayList<>(cur));
            return;
        }
        
        // 得到下一个要枚举的单词的前缀应该是什么
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < pos; i++) {
            sb.append(cur.get(i).charAt(pos));
        }
        
        // 得到以该前缀开头的所有单词
        List<String> nextWord = trie.getPrefix(sb.toString());
        for (int i = 0; i < nextWord.size(); i++) {
        	// 判断一下，如果枚举了nextWord[i]这个单词的话，后面单词应该拥有的前缀是否存在对应的单词，
        	// 如果不存在，那直接不要枚举当前单词
            if (!checkPref(pos, nextWord.get(i), cur, words[0].length(), trie)) {
                continue;
            }
            
            // 枚举当前单词作为下一个单词
            cur.add(nextWord.get(i));
            dfs(pos + 1, cur, words, res, trie);
            cur.remove(cur.size() - 1);
        }
    }
    
    private boolean checkPref(int pos, String next, List<String> cur, int wordLen, Trie trie) {
    	// 算一下pos之后的列的竖着看的前缀
        for (int i = pos + 1; i < wordLen; i++) {
            StringBuilder sb = new StringBuilder();
            for (int j = 0; j < pos; j++) {
                sb.append(cur.get(j).charAt(i));
            }
            sb.append(next.charAt(i));
            
            // 如果没有单词以sb为前缀，那返回false
            if (trie.getPrefix(sb.toString()).isEmpty()) {
                return false;
            }
        }
        
        return true;
    }
}

时空复杂度依赖于具体有哪些单词。