1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
|
class Trie:
def __init__(self):
"""
Initialize your data structure here.
"""
self.children = dict()
self.isEnd = False
def insert(self, word: str) -> None:
"""
Inserts a word into the trie.
"""
cur = self
for ch in word:
if ch not in cur.children:
cur.children[ch] = Trie()
cur = cur.children[ch]
cur.isEnd = True
def search(self, word: str):
"""
Returns if the word is in the trie.
"""
cur = self
res = []
for ch in word:
if ch not in cur.children:
return res
cur = cur.children[ch]
cur.dfs(word, res)
return res
def dfs(self, word, res):
cur = self
if cur and cur.isEnd:
res.append(word)
return
for ch in cur.children:
node = cur.children[ch]
node.dfs(word + ch, res)
class Solution:
def backtrace(self, index, size, path, res, trie_tree):
if index == size:
res.append(path[:])
return
next_prefix = "" # 下一行的前缀
for i in range(index):
next_prefix += path[i][index]
next_words_with_prefix = trie_tree.search(next_prefix)
for word in next_words_with_prefix:
path.append(word)
self.backtrace(index + 1, size, path, res, trie_tree)
path.pop(-1)
def wordSquares(self, words: List[str]) -> List[List[str]]:
trie_tree = Trie()
for word in words:
trie_tree.insert(word)
size = len(words[0])
res = []
path = []
for word in words:
path.append(word)
self.backtrace(1, size, path, res, trie_tree)
path.pop(-1)
return res
|