題目網址：https://leetcode.cn/problems/word-squares/

題意：給一不重複的 string set words, 返回其中所有的 word square。words 中的同一個單字可以被多次使用, 以任意順序返回答案。

word square 的定義：從第 k 列和第 k 行來看都是相同的 string, 其中 0 ≤ k ≤ max(numRows, numColumns)。

e.g. words = ["ball","area","lead","lady"] 為 word square, 因為每個單字從水平和垂直方向上來看都是相同的。

Solution：

想法：利用 Trie + DFS + Backtracking, 逐列填入單字, 新的一列填入的單字可由前面的單字得到, 其中第 idx 列的單字之 prefix 可由前面每個單字的第 idx + 1 個字母所串連

e.g. words = ["wall","able","area","lead","lady"]

  • 從 wall 開始, 作為第一列的單字

  • 根據 word square 的對稱性, 我們知道第二列的單字應以 a 為 prefix 開頭(因為第一列第二行的字母為 a)

  • 在 words 中, 有兩個 prefix 為 a 的單字(即 able, area)。因此下一步我們要對兩個單字依次進行嘗試

  • 選擇 able 作為第二列的單字。根據對稱性, 我們知道第三列的元素應以 ll 作為開頭。但是, words 中並沒有這樣的單字, 因此我們要放棄此次的嘗試, 並回到上一個狀態(第一列已填入)

  • 選擇 area 作為第二列的單字。根據對稱性, 我們知道第三列的元素應以 le 作為開頭。這次我們可以在 words 找到符合的單字, 即 lead

  • 選擇 lead 作為第三列的單字。根據對稱性, 我們知道第四列的元素應以 lad 作為開頭, 我們可以在 words 找到符合的單字, 即 lady

  • 得到 ["ball","area","lead","lady"] 一組解

class TrieNode{
public:
    TrieNode* children[26] = {nullptr};
    int idx = -1; // end node 紀錄該 word 在 words 中的 idx
};

class Trie{
public:
    TrieNode *root;

    Trie() : root(new TrieNode()) {}

    void insert(string& word, int& idx){
        TrieNode *node = root;

        for (const auto& ch : word) {
            const int i = ch - 'a';
            if (!node->children[i]) {
                node->children[i] = new TrieNode();
            }
            node = node->children[i];
        }

        node->idx = idx; // end node 之 idx 為非 -1 之值
    }

    // 返回 prefix 的 end node
    TrieNode* search(const string& prefix){
        TrieNode *node = root;

        for (const auto& ch : prefix) {
            const int i = ch - 'a';
            if (!node->children[i]) {
                return nullptr;
            }
            node = node->children[i];
        }

        return node;
    }

    // 取得以 node 為 root 的 subtree 中所有的 word(也就是以 prefix 為開頭的所有 word)
    void get(const TrieNode* node, vector<string>& words, vector<string>& prefixWords){
        if (!node) {
            return;
        }

        // 遇到 end node of a word, 則加入到 prefixWords 中, 並繼續往下遍歷
        if (node->idx != -1) {
            prefixWords.emplace_back(words[node->idx]);
        }

        // 往下遍歷
        for (int i = 0; i < 26; ++i) {
            get(node->children[i], words, prefixWords);
        }
    }
};

class Solution {
public:
    vector<vector<string>> wordSquares(vector<string>& words) {
        n = words[0].size();

        // build trie
        for (int i = 0; i < words.size(); ++i) {
            trie.insert(words[i], i);
        }

        // 遍歷每個單字為第0列的情況
        for (const auto& word : words) {
            cur.emplace_back(word);
            dfs(0, words);
            cur.pop_back();
        }

        return res;
    }

private:
    int n;
    Trie trie;
    vector<string> cur; // 當前 word square 填入的情況
    vector<vector<string>> res;

    // 代表第 idx 列已被填入
    void dfs(const int idx, vector<string>& words){
        // 代表第 0 ~ (n - 1) 列都可以填入, 找到一組解
        if (idx == n - 1) {
            res.emplace_back(cur);
            return;
        }

        // 得到下一列應填入單字的 prefix
        string prefix;
        for (int i = 0; i <= idx; ++i) {
            prefix.push_back(cur[i][idx + 1]);
        }

        // 找到 prefix 之 end node(也就是所有以 prefix 開頭的 word 之 root)
        const auto node = trie.search(prefix);

        // 找到所有以 prefix 開頭的 word
        vector<string> prefixWords;
        trie.get(node, words, prefixWords);

        // 在下一列嘗試所有符合的 word
        for (const auto& word : prefixWords) {
            cur.emplace_back(word);
            dfs(idx + 1, words);
            cur.pop_back();
        }
    }
};
  • time：$O(n \cdot L + n \cdot 26 ^ L)$ ➔ n 是 words 中單字的個數, L 是單字的平均長度
    • $O(n \cdot L)$：build trie
    • $O(n \cdot 26 ^ L)$：dfs(0, words[i]) 以 words[i] 做第一列的單字嘗試所有的可能需花 $(26^L)$, 因為 trie 中每個 node 有 26 個分支, 所以 trie 中最多有 $O(26^L)$ 個 node 要嘗試。總共有 $n$ 個 words[i] 要嘗試作為第一列的單字。
  • space：$O(n \cdot L)$ ➔ trie, worse case：trie 中的單字都沒有重複的 prefix