https://leetcode.com/problems/word-search-ii/

The following solution get TLE.

class Solution:
def findWords(self, matrix: List[List[str]], words: List[str]) -> List[str]:
if not matrix or not matrix[0]:return []
res = []
R, C = len(matrix), len(matrix[0])
d = collections.defaultdict(set)
for i in range(R):
for j in range(C):
d[matrix[i][j]].add((i,j))
def dfs(word, seen, i, j):
if len(word)==1:return True
ch = word[1]
neighbords = []
for di, dj in [(1,0),(-1,0),(0,1),(0,-1)]:
r, c = i+di, j+dj
if 0<=r<R and 0<=c<C and matrix[r][c]==ch and (r,c) not in seen:
new_seen = seen.copy()
new_seen.add((r,c))
if dfs(word[1:], new_seen, r,c):
return True
return False
for word in words:
if any(dfs(word, {(i,j)},i,j) for i,j in d[word[0]]):
res.append(word)
return res

Alternative backtracking solution

Change the code to a backtracking solution, we still get TLE

class Solution:
def findWords(self, matrix: List[List[str]], words: List[str]) -> List[str]:
if not matrix or not matrix[0]:return []
res = []
R, C = len(matrix), len(matrix[0])
d = collections.defaultdict(set)
for i in range(R):
for j in range(C):
d[matrix[i][j]].add((i,j))
def dfs(word, seen, i, j):
if len(word)==1:return True
ch = word[1]
neighbords = []
for di, dj in [(1,0),(-1,0),(0,1),(0,-1)]:
r, c = i+di, j+dj
if 0<=r<R and 0<=c<C and matrix[r][c]==ch and (r,c) not in seen:
seen.add((r,c))
if dfs(word[1:], seen, r,c):
return True
seen.discard((r,c))
return False
for word in words:
if any(dfs(word, {(i,j)},i,j) for i,j in d[word[0]]):
res.append(word)
return res

A cleaner backtracking solution (still get TLE. C++ version can get AC)

class Solution:
def findWords(self, matrix: List[List[str]], words: List[str]) -> List[str]:
if not matrix or not matrix[0]:return []
res = []
R, C = len(matrix), len(matrix[0])
d = collections.defaultdict(set)
for i in range(R):
for j in range(C):
d[matrix[i][j]].add((i,j))
def dfs(word, k, r, c):
if r<0 or r>=R or c<0 or c>=C or matrix[r][c]!=word[k]:return False
if k==len(word)-1:return True
cur = matrix[r][c]
matrix[r][c] = '#'
found = dfs(word, k+1, r+1,c) or dfs(word, k+1, r-1,c) or dfs(word, k+1, r,c+1) or dfs(word, k+1, r,c-1)
matrix[r][c]=cur
return found

for word in words:
for i,j in d[word[0]]:
if dfs(word, 0, i, j):
res.append(word)
break
return res

Trie solution get AC

_END = "_END"
class Trie:
def __init__(self):
self.trie = {}
def insert(self, word):
current_d = self.trie
for w in word:
current_d = current_d.setdefault(w, {})
current_d[_END] = _END
class Solution:
def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
if not board or not board[0]:return []
R, C =len(board), len(board[0])
res, trie = set(), Trie()
def dfs(i,j, word, current_d):
if _END in current_d:
res.add(word)

if i<0 or i>=R or j<0 or j>=C or board[i][j] not in current_d:return
# backtracking
temp = board[i][j]
current_d = current_d[temp]
board[i][j] = "#"
for di, dj in [(1,0),(-1,0),(0,1),(0,-1)]:
r, c = i+di, j+dj
dfs(r,c,word+temp,current_d)
board[i][j] = temp

for word in words:
trie.insert(word)
cur_d = trie.trie
for i in range(R):
for j in range(C):
dfs(i, j, "", cur_d)
return list(res)

Article about Trie can be found here.

Data Scientist/MLE/SWE @takemobi