Question

425. Word Squares

中文文档

Description

Given an array of unique strings words, return _all the _word squares_ you can build from _words. The same word from words can be used multiple times. You can return the answer in any order.

A sequence of strings forms a valid word square if the kth row and column read the 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.

 

Example 1:

Input: words = ["area","lead","wall","lady","ball"]
Output: [["ball","area","lead","lady"],["wall","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: words = ["abat","baba","atan","atal"]
Output: [["baba","abat","baba","atal"],["baba","abat","baba","atan"]]
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).

 

Constraints:

• 1 <= words.length <= 1000
• 1 <= words[i].length <= 4
• All words[i] have the same length.
• words[i] consists of only lowercase English letters.
• All words[i] are unique.

Solutions

Solution 1

class Trie:
  def __init__(self):
    self.children = [None] * 26
    self.v = []

  def insert(self, w, i):
    node = self
    for c in w:
      idx = ord(c) - ord('a')
      if node.children[idx] is None:
        node.children[idx] = Trie()
      node = node.children[idx]
      node.v.append(i)

  def search(self, w):
    node = self
    for c in w:
      idx = ord(c) - ord('a')
      if node.children[idx] is None:
        return []
      node = node.children[idx]
    return node.v


class Solution:
  def wordSquares(self, words: List[str]) -> List[List[str]]:
    def dfs(t):
      if len(t) == len(words[0]):
        ans.append(t[:])
        return
      idx = len(t)
      pref = [v[idx] for v in t]
      indexes = trie.search(''.join(pref))
      for i in indexes:
        t.append(words[i])
        dfs(t)
        t.pop()

    trie = Trie()
    ans = []
    for i, w in enumerate(words):
      trie.insert(w, i)
    for w in words:
      dfs([w])
    return ans

class Trie {
  Trie[] children = new Trie[26];
  List<Integer> v = new ArrayList<>();

  void insert(String word, int i) {
    Trie node = this;
    for (char c : word.toCharArray()) {
      c -= 'a';
      if (node.children[c] == null) {
        node.children[c] = new Trie();
      }
      node = node.children[c];
      node.v.add(i);
    }
  }

  List<Integer> search(String pref) {
    Trie node = this;
    for (char c : pref.toCharArray()) {
      c -= 'a';
      if (node.children[c] == null) {
        return Collections.emptyList();
      }
      node = node.children[c];
    }
    return node.v;
  }
}

class Solution {
  private Trie trie = new Trie();
  private String[] words;
  private List<List<String>> ans = new ArrayList<>();

  public List<List<String>> wordSquares(String[] words) {
    this.words = words;
    for (int i = 0; i < words.length; ++i) {
      trie.insert(words[i], i);
    }

    List<String> t = new ArrayList<>();
    for (String w : words) {
      t.add(w);
      dfs(t);
      t.remove(t.size() - 1);
    }
    return ans;
  }

  private void dfs(List<String> t) {
    if (t.size() == words[0].length()) {
      ans.add(new ArrayList<>(t));
      return;
    }
    int idx = t.size();
    StringBuilder pref = new StringBuilder();
    for (String x : t) {
      pref.append(x.charAt(idx));
    }
    List<Integer> indexes = trie.search(pref.toString());
    for (int i : indexes) {
      t.add(words[i]);
      dfs(t);
      t.remove(t.size() - 1);
    }
  }
}

type Trie struct {
  children [26]*Trie
  v    []int
}

func newTrie() *Trie {
  return &Trie{}
}
func (this *Trie) insert(word string, i int) {
  node := this
  for _, c := range word {
    c -= 'a'
    if node.children[c] == nil {
      node.children[c] = newTrie()
    }
    node = node.children[c]
    node.v = append(node.v, i)
  }
}
func (this *Trie) search(word string) []int {
  node := this
  for _, c := range word {
    c -= 'a'
    if node.children[c] == nil {
      return []int{}
    }
    node = node.children[c]
  }
  return node.v
}

func wordSquares(words []string) [][]string {
  trie := newTrie()
  for i, w := range words {
    trie.insert(w, i)
  }
  ans := [][]string{}
  var dfs func([]string)
  dfs = func(t []string) {
    if len(t) == len(words[0]) {
      cp := make([]string, len(t))
      copy(cp, t)
      ans = append(ans, cp)
      return
    }
    idx := len(t)
    pref := []byte{}
    for _, v := range t {
      pref = append(pref, v[idx])
    }
    indexes := trie.search(string(pref))
    for _, i := range indexes {
      t = append(t, words[i])
      dfs(t)
      t = t[:len(t)-1]
    }
  }
  for _, w := range words {
    dfs([]string{w})
  }
  return ans
}