Question

425. 单词方块

English Version

题目描述

给定一个单词集合 words （没有重复），找出并返回其中所有的 单词方块 。 words 中的同一个单词可以被 多次 使用。你可以按 任意顺序 返回答案。

一个单词序列形成了一个有效的 单词方块 的意思是指从第 k 行和第 k 列  0 <= k < max(numRows, numColumns) 来看都是相同的字符串。

• 例如，单词序列 ["ball","area","lead","lady"] 形成了一个单词方块，因为每个单词从水平方向看和从竖直方向看都是相同的。

 

示例 1：

输入：words = ["area","lead","wall","lady","ball"]
输出: [["ball","area","lead","lady"],
["wall","area","lead","lady"]]
解释：
输出包含两个单词方块，输出的顺序不重要，只需要保证每个单词方块内的单词顺序正确即可。 

示例 2：

输入：words = ["abat","baba","atan","atal"]
输出：[["baba","abat","baba","atal"],
["baba","abat","baba","atan"]]
解释：
输出包含两个单词方块，输出的顺序不重要，只需要保证每个单词方块内的单词顺序正确即可。 

 

提示:

• 1 <= words.length <= 1000
• 1 <= words[i].length <= 4
• words[i] 长度相同
• words[i] 只由小写英文字母组成
• words[i] 都 各不相同

解法

方法一：前缀树 + DFS

根据已添加单词确定下一个单词的前缀，继续进行搜索。

比如已经添加了两个单词 $ball$ 和 $area$，要添加下一个单词，我们首先要获取下一个单词的前缀，第一个字母是第一个单词的第三个位置 $l$，第二个字母是第二个单词的第三个位置 $e$，这样就构成了前缀 $le$。然后找出所有前缀为 $le$ 的单词，作为下一个单词。

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
}