Word Squares

hard trie backtracking string

Problem

Given a set of equal-length words, return all the word squares you can build using each word at most once.

Inputwords = ["area","lead","wall","lady","ball"]

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

Each row also reads as the corresponding column.

words (comma-separated)

def word_squares(words):
    if not words: return []
    n = len(words[0])
    pref = {}
    for w in words:
        for k in range(n + 1):
            pref.setdefault(w[:k], []).append(w)
    res, cur = [], []
    def back():
        if len(cur) == n:
            res.append(cur[:]); return
        k = len(cur)
        p = "".join(row[k] for row in cur)
        for w in pref.get(p, []):
            cur.append(w); back(); cur.pop()
    back()
    return res

function wordSquares(words) {
  if (!words.length) return [];
  const n = words[0].length;
  const pref = new Map();
  for (const w of words) {
    for (let k = 0; k <= n; k++) {
      const p = w.slice(0, k);
      if (!pref.has(p)) pref.set(p, []);
      pref.get(p).push(w);
    }
  }
  const res = [], cur = [];
  function back() {
    if (cur.length === n) { res.push(cur.slice()); return; }
    const k = cur.length;
    let p = "";
    for (const row of cur) p += row[k];
    for (const w of (pref.get(p) || [])) {
      cur.push(w); back(); cur.pop();
    }
  }
  back();
  return res;
}

class Solution {
    Map<String, List<String>> pref = new HashMap<>();
    int n;
    List<List<String>> res = new ArrayList<>();
    Deque<String> cur = new ArrayDeque<>();
    public List<List<String>> wordSquares(String[] words) {
        if (words.length == 0) return res;
        n = words[0].length();
        for (String w : words)
            for (int k = 0; k <= n; k++)
                pref.computeIfAbsent(w.substring(0, k), x -> new ArrayList<>()).add(w);
        back();
        return res;
    }
    void back() {
        if (cur.size() == n) { res.add(new ArrayList<>(cur)); return; }
        int k = cur.size();
        StringBuilder p = new StringBuilder();
        for (String row : cur) p.append(row.charAt(k));
        for (String w : pref.getOrDefault(p.toString(), Collections.emptyList())) {
            cur.addLast(w); back(); cur.removeLast();
        }
    }
}

class Solution {
    unordered_map<string, vector<string>> pref;
    int n;
    vector<vector<string>> res;
    vector<string> cur;
    void back() {
        if ((int)cur.size() == n) { res.push_back(cur); return; }
        int k = cur.size();
        string p;
        for (auto& row : cur) p += row[k];
        for (auto& w : pref[p]) { cur.push_back(w); back(); cur.pop_back(); }
    }
public:
    vector<vector<string>> wordSquares(vector<string>& words) {
        if (words.empty()) return res;
        n = words[0].size();
        for (auto& w : words)
            for (int k = 0; k <= n; k++)
                pref[w.substr(0, k)].push_back(w);
        back();
        return res;
    }
};

Explanation

In a word square, the k-th row must read the same as the k-th column. The trick is to build the square row by row: once some rows are placed, the next row's starting letters are forced by the columns already formed.

To find candidates quickly we index words by every prefix: pref[w[:k]] lists all words starting with that prefix. This is built once up front.

The backtracking function back() looks at how many rows are filled, say k, then computes the required prefix p by reading down column k of the current rows (row[k] for each row). Only words in pref.get(p, []) can legally go next.

For each valid word we push it, recurse, then pop it (undo) and try the next. When the number of rows equals n, we have a full square and save a copy.

Example: with the first row wall, the second row must start with a (column 0 already spells w,a,…), so the prefix lookup hands us area and we continue. Forcing prefixes this way prunes the search dramatically.

Time: O(N · L · 26) Space: O(N · L)