425. Word Squares

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).

Rust Solution

struct Solution;

use std::collections::HashMap;

impl Solution {
    fn word_squares(words: Vec<String>) -> Vec<Vec<String>> {
        let n = words.len();
        let m = words[0].len();
        let words: Vec<Vec<char>> = words.into_iter().map(|s| s.chars().collect()).collect();
        let mut indexes: HashMap<Vec<char>, Vec<&[char]>> = HashMap::new();
        for word in words.iter() {
            for end in 0..=m {
                indexes.entry(word[..end].to_vec()).or_default().push(word);
            }
        }
        let mut cur = vec![];
        let mut res = vec![];
        Self::dfs(0, &mut cur, &mut res, &indexes, m, n);
        res
    }

    fn dfs<'a>(
        start: usize,
        cur: &mut Vec<&'a [char]>,
        res: &mut Vec<Vec<String>>,
        indexes: &HashMap<Vec<char>, Vec<&'a [char]>>,
        m: usize,
        n: usize,
    ) {
        if start == m {
            let square = (0..m).map(|i| cur[i].iter().collect()).collect();
            res.push(square);
        } else {
            let mut prefix = vec![];
            for i in 0..start {
                prefix.push(cur[i][start]);
            }
            if let Some(candidates) = indexes.get(&prefix) {
                for word in candidates {
                    cur.push(word);
                    Self::dfs(start + 1, cur, res, indexes, m, n);
                    cur.pop();
                }
            }
        }
    }
}

#[test]
fn test() {
    let words = vec_string!["area", "lead", "wall", "lady", "ball"];
    let res = vec_vec_string![
        ["wall", "area", "lead", "lady"],
        ["ball", "area", "lead", "lady"]
    ];
    assert_eq!(Solution::word_squares(words), res);
    let words = vec_string!["abat", "baba", "atan", "atal"];
    let res = vec_vec_string![
        ["baba", "abat", "baba", "atan"],
        ["baba", "abat", "baba", "atal"]
    ];
    assert_eq!(Solution::word_squares(words), res);
}

Having problems with this solution? Click here to submit an issue on github.