Word search

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.