golang 解法, 执行用时: 244 ms, 内存消耗: 32.4 MB, 提交时间: 2023-10-22 08:51:35

func wordSquares(words []string) [][]string {
    if len(words)==0 || len(words[0])==0{
        return [][]string{}
    }
    solutions:=[][]string{}
    board:=make([]string, len(words[0]))
    dic:=map[string][]string{} // 行前缀哈希表
    for _,i:=range words{ // 添加行前缀对应的字符串
        for j:=0;j<len(i);j++{
            dic[i[:j]] = append(dic[i[:j]], i)
        }
    }
    var backtrack func(int)
    backtrack = func(cnt int){
        if cnt==len(words[0]){ // 回溯结束条件
            solutions = append(solutions, append([]string(nil),board...))
            return
        }
        if _,ok:=dic[board[cnt]];ok{ // 如果现在的行前缀没有对应的字符串，则不用继续回溯下去
            for _,i:=range dic[board[cnt]]{
                board[cnt] = i
                for j:=cnt+1;j<len(words[0]);j++{
                    board[j] = board[j]+string(board[cnt][j])
                }
                backtrack(cnt+1)
                for j:=cnt+1;j<len(words[0]);j++{
                    board[j] = board[j][:len(board[j])-1]
                }
                board[cnt] = board[cnt][:cnt]
            }
        }
    }
    backtrack(0)
    return solutions
}

python3 解法, 执行用时: 464 ms, 内存消耗: 42.2 MB, 提交时间: 2023-10-22 08:51:11

class Trie:
    def __init__(self):
        self.child = dict()
        self.isWord = False

    def insert(self, word: str) -> None:    
        root = self
        for c in word:
            if c not in root.child:
                root.child[c] = Trie()
            root = root.child[c]
        root.isWord = True
    
    def starts_with(self, prefix: str) -> bool:
        root = self
        for c in prefix:
            if c not in root.child:
                return False
            root = root.child[c]
        return True
    
    def is_word(self, word: str) -> bool:
        root = self
        for c in word:
            if c not in root.child:
                return False
            root = root.child[c]
        return root.isWord

    def query(self, prefix: str) -> List[str]:
        res = []
        root = self
        for c in prefix:
            if c not in root.child:
                return []
            root = root.child[c]

        def dfs(prefix: str,  root) -> None:
            if root.isWord == True:
                res.append(prefix)
            for c in root.child:
                dfs(prefix + c, root.child[c])

        dfs(prefix, root)
        return res

class Solution:
    def wordSquares(self, words: List[str]) -> List[List[str]]:
        #挑选单词时顺序是随意的
        #关于主对角线对称 当前已有的行们决定下一行的前缀

        #------ 先建trie树
        T = Trie()
        for word in words:
            T.insert(word)
        
        #------------------- 回溯
        def backtrace(idx: int) -> None:
            if idx == wordLen:          #剪枝
                res.append(path[:])
                return
            nxt_prefix = ""             #下一行的前缀
            for i in range(idx):
                nxt_prefix += path[i][idx]

            nxt_words_with_prefix = T.query(nxt_prefix)
            for word in nxt_words_with_prefix:
                path.append(word)
                backtrace(idx + 1)
                path.pop(-1)

        #---------------- 运行回溯算法
        wordLen = len(words[0])
        res = []

        path = []
        for word in words:
            path.append(word)
            backtrace(1)
            path.pop(-1)
        return res

cpp 解法, 执行用时: 128 ms, 内存消耗: 57.5 MB, 提交时间: 2023-10-22 08:50:49

class Solution {
public:
    struct TrieNode {
        vector<TrieNode*> children;
        vector<int> indexs;
        TrieNode() : children(26, nullptr) { }
    };
    
    TrieNode* buildTrie(const vector<string> &words) {
        TrieNode *root = new TrieNode();
        for (int i = 0; i < words.size(); ++i) {
            TrieNode *p = root;
            for (const auto &ch : words[i]) {
                int j = ch - 'a';
                if (p->children[j] ==  nullptr) {
                    p->children[j] = new TrieNode();
                }
                p = p->children[j];
                p->indexs.push_back(i);
            }
        }
        return root;
    }

    vector<vector<string>> wordSquares(vector<string>& words) {
        vector<vector<string>> res;        
        int size = words[0].size();
        TrieNode *root = buildTrie(words);
        vector<string> out(size);
        for (const auto &word : words) {
            out[0] = word;
            helper(words, root, 1, out, res);
        }
        return res;
    }

    void helper(vector<string>& words, TrieNode *root, int level, vector<string> &out, vector<vector<string>> &res) {
        if (level >= words[0].size()) {
            res.push_back(out);
            return;
        }
        string str;
        for (int i = 0; i < level; ++i) {
            str += out[i][level];
        }
        TrieNode *t = root;
        for (int i = 0; i < str.size(); ++i) {
            if (!t->children[str[i] - 'a']) return;
            t = t->children[str[i] - 'a'];
        }
        for (int idx : t->indexs) {
            out[level] = words[idx];
            helper(words, root, level + 1, out, res);
        }
    }
};

cpp 解法, 执行用时: 244 ms, 内存消耗: 100.3 MB, 提交时间: 2023-10-22 08:50:35

class Solution {
public:
    vector<vector<string>> wordSquares(vector<string>& words) {
        vector<vector<string>> res;
        unordered_map<string, set<string>> prefix_word_map;
        int size = words[0].size();
        for (auto &word : words) {
            for (int i = 0; i < size; ++i) {
                string prefix = word.substr(0, i);
                prefix_word_map[prefix].insert(word);
            }
        }
        vector<vector<char>> matrix(size, vector<char>(size));
        helper(0, size, matrix, prefix_word_map, res);
        return res;
    }

    void helper(int i, int size, vector<vector<char>> &matrix, unordered_map<string, set<string>> &prefix_word_map, vector<vector<string>> &res) {
        if (i == size) {            
            vector<string> tmp;
            for (int j = 0; j < size; ++j) {
                tmp.push_back(string(matrix[j].begin(), matrix[j].end()));                
            }
            res.push_back(tmp);
            return;
        }
        string prefix = string(matrix[i].begin(), matrix[i].begin() + i);
        for (string word : prefix_word_map[prefix]) {
            matrix[i][i] = word[i];
            int j = i + 1;
            for (; j < size; ++j) {
                matrix[i][j] = word[j];//填充行
                matrix[j][i] = word[j];//填充列
                if (prefix_word_map.count(string(matrix[j].begin(), matrix[j].begin() + i + 1)) == 0) break;
            }
            if (j == size) helper(i + 1, size, matrix, prefix_word_map, res);
        }
    }
};

cpp 解法, 执行用时: 108 ms, 内存消耗: 57 MB, 提交时间: 2023-10-22 08:50:27

class TrieNode {
public:
    string* word = nullptr;
    TrieNode* children[26];
    
    TrieNode() {
        for (int i = 0; i < 26; ++i)
            children[i] = nullptr;
    }
    
    ~TrieNode() {
        for (int i = 0; i < 26; ++i)
            delete children[i];
    }
};


class Trie {
private:
    TrieNode rootNode;
public:
    TrieNode* root;
    
    Trie() {
        root = &rootNode;
    }
    
    void insert(string& word) {
        TrieNode* node = root;
        for (char c: word) {
            int index = c - 'a';
            if (node->children[index] == nullptr)
                node->children[index] = new TrieNode;
            node = node->children[index];
        }
        node->word = &word;
    }
};


class Solution {
private:
    Trie trie;
    void helper(TrieNode* node, 
                int index,
                vector<TrieNode*>& columns, 
                vector<string>& candidate, 
                vector<vector<string>>& result) {
        
        if (node->word != nullptr) {
            candidate.push_back(*node->word);
            if (candidate.size() == columns.size())
                result.push_back(candidate);
            else
                helper(columns[candidate.size()], candidate.size(), columns, candidate, result);
            candidate.pop_back();
            return;
        }
        
        for (int i = 0; i < 26; ++i) {
            if (node->children[i] != nullptr && columns[index]->children[i] != nullptr) {
                TrieNode* parent = columns[index];
                columns[index] = columns[index]->children[i];
                helper(node->children[i], index+1, columns, candidate, result);
                columns[index] = parent;
            }
        }
    }
    
public:
    vector<vector<string>> wordSquares(vector<string>& words) {
        for (string& word: words)
            trie.insert(word);
        
        int n = words[0].size();
        vector<TrieNode*> columns(n, trie.root);
        vector<string> candidate;
        vector<vector<string>> result;
        helper(trie.root, 0, columns, candidate, result);
        
        return result;
    }
};

java 解法, 执行用时: 28 ms, 内存消耗: 66.2 MB, 提交时间: 2023-10-22 08:49:48

class Solution {
    public List<List<String>> wordSquares(String[] words) {
        int n = words[0].length();
        Trie trie = new Trie();
        for (String word: words)
            trie.insert(word, n);
        
        List<List<String>> result = new LinkedList();
        String[] finished = new String[n];
        char[][] board = new char[n][n];
        
        search(trie.root, trie.root, board, 0, 0, n, finished, result);
        
        return result;
    }
    
    /*
    *   current：当前搜索节点
    *   root：根节点
    *   board：字符矩阵
    *   i：正要处理的行（第几个单词）
    *   j：正要处理的列（第几个字符）
    *   n：单词长度（矩阵宽度）
    *   finished：已经填好的单词（0 ~ i-1 已经填好）
    *   result：输出结果
    */
    private void search(TrieNode current, TrieNode root, char[][] board, int i, int j, int n, String[] finished, List<List<String>> result) {
        
        // 全部n个单词填好，输出结果
        if (i == n) {
            result.add(new ArrayList<String>(Arrays.asList(finished)));
            return;
        }
        
        // 第i个单词填好，放入finished
        if (j == n) {
            finished[i] = current.word;
            search(root, root, board, i+1, 0, n, finished, result);
            return;
        }
        
        // 之前已经填过的格子，检查是否可以找到
        // 比如在填第一个单词wall时，(0, 1) 处字符填了a，同时在(1, 0)处也填了a
        // 那么在处理第二个单词时，先要看第一个字符（也就是 i = 0, j = 1 处是否可以为a
        // 如果不能为 a，结束该搜索过程，回溯
        if (j < i) {
            TrieNode child = current.getChild(board[i][j]);
            if (child != null)
                search(child, root, board, i, j+1, n, finished, result);
            return;
        }
        
        // 尝试在该点添堵所有可能的字符
        for (int c = 0; c < 26; ++c) {
            if (current.children[c] != null) {
                // 如果在填第一个单词时，同时可以判断每个字符是否可以作为起始字符，从而减少递归分支
                if (i == 0 && root.children[c] == null)
                    continue;
                board[j][i] = board[i][j] = (char) ('a' + c);
                search(current.children[c], root, board, i, j+1, n, finished, result);
            }
        }
    }
    
}

class Trie {
    public TrieNode root = new TrieNode();
    
    public void insert(String word, int n) {
        TrieNode node = root;
        int charNo;
        for (int i = 0; i < n; ++i) {
            charNo = word.charAt(i) - 'a';
            if (node.children[charNo] == null)
                node.children[charNo] = new TrieNode();
            node = node.children[charNo];
        }
        node.word = word;
    }
}

class TrieNode {
    public TrieNode[] children = new TrieNode[26];
    public String word = null;
    
    public TrieNode getChild(char c) {
        return children[c - 'a'];
    }
}