题目

给定一个单词集合（没有重复），找出其中所有的 单词方块。

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

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

  1. b a l l
  2. a r e a
  3. l e a d
  4. l a d y

注意：

  • 单词个数大于等于 1 且不超过 500。
  • 所有的单词长度都相同。
  • 单词长度大于等于 1 且不超过 5。
  • 每个单词只包含小写英文字母 a-z。

示例 1：

输入：
["area","lead","wall","lady","ball"]

输出：
[
  [ "wall",
    "area",
    "lead",
    "lady"
  ],
  [ "ball",
    "area",
    "lead",
    "lady"
  ]
]

解释：
输出包含两个单词方块，输出的顺序不重要，只需要保证每个单词方块内的单词顺序正确即可。

示例 2：

输入：
["abat","baba","atan","atal"]

输出：
[
  [ "baba",
    "abat",
    "baba",
    "atan"
  ],
  [ "baba",
    "abat",
    "baba",
    "atal"
  ]
]

解释：
输出包含两个单词方块，输出的顺序不重要，只需要保证每个单词方块内的单词顺序正确即可。

方案一

class Trie:

    def __init__(self, words=[]):
        self.lookup = {}
        for word in words:
            self.add(word)

    def add(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        node = self.lookup
        for c in word:
            if c not in node:
                node[c] = {}
            node = node[c]
        node['is_word'] = True

    def has(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        node = self.lookup
        for c in word:
            if c not in node:
                return False
            node = node[c]
        # 判断是否是单词
        if 'is_word' in node:
            return True
        return False

    def startsWith(self, prefix: str, limit: int) -> [str]:
        """
        Returns if there is any word in the trie that starts with the given prefix and length equal limit.
        需要保证 len(prefix) > limit，在不做特殊处理
        """
        node = self.lookup
        index = 0
        for c in prefix:
            index += 1
            if c not in node:
                return []
            node = node[c]
        return self._getAllWords(node, prefix, limit)

    def _getAllWords(self, node, prefix, limit):
        ret = []
        if len(prefix) > limit:
            return []
        if len(prefix) == limit and "is_word" in node:
            return [prefix]
        for key in node:
            ret.extend(self._getAllWords(node[key], prefix + key, limit))

        return ret


class Solution:
    def __init__(self):
        self.trie = Trie()

    def wordSquares(self, words: List[str]) -> List[List[str]]:
        if not words:
            return []

        for word in words:
            self.trie.add(word)

        # 比如有[ball]，我们知道单词方块是 4X4，换句话说每个单词长度都为4，其次下个单词是a字母开头的。
        # 假如我们找到一个a开头长度为4单词，变成[ball, area]，那么第三个单词的前缀是le长度为4单词，依次类推。
        ret = []

        def backtrack(item: [str]):
            '''每一个方块'''
            if len(item) == len(item[0]):
                ret.append(item)
                return
            # 获取当前的前缀
            prefix = ""
            for it in item:
                prefix += it[len(item)] # 因为是方块，所以 len(item) == 当前方块的高
            # 获取前缀是 prefix 长度为 len(item[0]) 的单词
            prefix_words = self.trie.startsWith(prefix, len(item[0]))
            for pw in prefix_words:
                backtrack(item + [pw])

        for word in words:
            backtrack([word])

        return ret

原文

https://leetcode-cn.com/explore/learn/card/trie/168/practical-application-ii/653/