Word squares

use std::collections::HashMap;

impl Solution {
  pub fn word_squares(words: Vec<String>) -> Vec<Vec<String>> {
    let n = words[0].len();
    // Build prefix -> list of words map
    let mut prefix_map: HashMap<String, Vec<usize>> = HashMap::new();
    for (i, w) in words.iter().enumerate() {
      for l in 0..=w.len() {
        prefix_map.entry(w[..l].to_string()).or_default().push(i);
      }
    }

    let mut result = Vec::new();
    let mut square: Vec<String> = Vec::with_capacity(n);

    Self::backtrack(&words, &prefix_map, n, &mut square, &mut result);
    result
  }

  fn backtrack(
    words: &[String],
    prefix_map: &HashMap<String, Vec<usize>>,
    n: usize,
    square: &mut Vec<String>,
    result: &mut Vec<Vec<String>>,
  ) {
    if square.len() == n {
      result.push(square.clone());
      return;
    }
    let step = square.len();
    // Build the required prefix for the next word: column `step` so far
    let prefix: String = square.iter().map(|w| w.as_bytes()[step] as char).collect();
    if let Some(candidates) = prefix_map.get(&prefix) {
      for &idx in candidates {
        square.push(words[idx].clone());
        Self::backtrack(words, prefix_map, n, square, result);
        square.pop();
      }
    }
  }
}