横竖相同的单词方阵

Given a set of words (without duplicates), find all word squares you can build from them.

A sequence of words forms a valid word square if the kth row and column read the exact same string, where 0 ≤ k < max(numRows, numColumns).

For example, the word sequence ["ball","area","lead","lady"] forms a word square because each word reads the same both horizontally and vertically.

b a l l
a r e a
l e a d
l a d y

Note:

  1. There are at least 1 and at most 1000 words.
  2. All words will have the exact same length.
  3. Word length is at least 1 and at most 5.
  4. Each word contains only lowercase English alphabet a-z.

Example 1:

Input:
["area","lead","wall","lady","ball"]

Output:
[
  [ "wall",
    "area",
    "lead",
    "lady"
  ],
  [ "ball",
    "area",
    "lead",
    "lady"
  ]
]

Explanation:
The output consists of two word squares. The order of output does not matter (just the order of words in each word square matters).

Example 2:

Input:
["abat","baba","atan","atal"]

Output:
[
  [ "baba",
    "abat",
    "baba",
    "atan"
  ],
  [ "baba",
    "abat",
    "baba",
    "atal"
  ]
]

Explanation:
The output consists of two word squares. The order of output does not matter (just the order of words in each word square matters).
vector<vector<string>> wordSquares(vector<string>& words) {
    // b a l l
    // a r e a
    // l e a d
    // l a d y
    // 第0行选定后,第1行要以a开头(ball[1]);
    // 第1行选定后,第2行要以le开头(ball[2]、area[2]);
    // 第2行选定后,第3行要以lad开头(ball[3]、area[3]、lead[3]);
    // ...        第k行要以串rows[0..k-1][k]开头
    // 需要根据前缀找单词的功能,可以用prefix=>wordIdx[]的哈希表,或者tire节点中存wordIdx[]
    if (words.empty()) return {};
    unordered_map<string, vector<int>> mp; // prefix=>wordIdx[]
    for (int i = 0; i < words.size(); i++) {
        for (int j = 0; j < words[i].size(); j++) {
            auto prefix = words[i].substr(0, j + 1);
            mp[prefix].push_back(i);
        }
    }
    // 回溯法
    vector<vector<string>> ans;
    vector<string> rows;
    for (auto &word : words) {
        rows.push_back(word);
        search(rows, words, mp, ans);
        rows.pop_back();
    }
    return ans;
}

void search(vector<string> &rows, vector<string> &words,
            unordered_map<string, vector<int>> &mp, vector<vector<string>> &ans) {
    if (rows.size() == words[0].size()) { // 有len(word)行
        ans.push_back(rows);
        return;
    }
    
    int k = rows.size(); // 第k行要以rows[0..k-1][k]开头
    string prefix;
    for (int i = 0; i < k; i++) {
        prefix += rows[i][k];            
    }
    for (int idx : mp[prefix]) { // 前缀为prefix的那些单词,可作为rows[]新一行的候选
        rows.push_back(words[idx]);
        search(rows, words, mp, ans);
        rows.pop_back();
    }
}