Merge pull request neetcode-gh#298 from niketanmoon/main

Solved TLE issue. Added removeWord function to update the Trie after …

212-Word-Search-II.py

@@ -2,64 +2,65 @@ class TrieNode:
  def __init__(self):
  self.children = {}
  self.isWord = False
- self.refs = 0
+
+class Trie:
+ def __init__(self):
+ self.root = TrieNode()
 
  def addWord(self, word):
- cur = self
- cur.refs += 1
+ current = self.root
  for c in word:
- if c not in cur.children:
- cur.children[c] = TrieNode()
- cur = cur.children[c]
- cur.refs += 1
- cur.isWord = True
+ if c not in current.children:
+ current.children[c] = TrieNode()
+ current = current.children[c]
+ current.isWord = True
 
  def removeWord(self, word):
- cur = self
- cur.refs -= 1
- for c in word:
- if c in cur.children:
-  cur = cur.children[c]
-  cur.refs -= 1
+ node = self.root
+ path = [(None, node)]
+ for char in word:
+ node = node.children[char]
+ path.append((char, node))
+ node.isWord = False
 
+ for i in range(len(path) - 1, 0, -1):
+ char, currNode = path[i][0], path[i][1]
+ parentNode = path[i - 1][1]
+ if currNode.isWord or currNode.children:
+ break
+ del parentNode.children[char]
 
 class Solution:
  def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
- root = TrieNode()
- for w in words:
- root.addWord(w)
+ trie = Trie()
 
- ROWS, COLS = len(board), len(board[0])
- res, visit = set(), set()
+ # Converting words list to a trie
+ for word in words:
+ trie.addWord(word)
+
+ rows, columns = len(board), len(board[0])
+ result, visit = set(), set()
 
  def dfs(r, c, node, word):
- if (
- r < 0
- or c < 0
- or r == ROWS
- or c == COLS
- or board[r][c] not in node.children
- or node.children[board[r][c]].refs < 1
- or (r, c) in visit
- ):
+ if (r < 0 or c < 0 or
+ r == rows or c == columns or
+ board[r][c] not in node.children or (r, c) in visit):
  return
 
  visit.add((r, c))
  node = node.children[board[r][c]]
  word += board[r][c]
  if node.isWord:
- node.isWord = False
- res.add(word)
- root.removeWord(word)
+ result.add(word)
+ trie.removeWord(word)
 
  dfs(r + 1, c, node, word)
  dfs(r - 1, c, node, word)
  dfs(r, c + 1, node, word)
  dfs(r, c - 1, node, word)
  visit.remove((r, c))
 
- for r in range(ROWS):
- for c in range(COLS):
- dfs(r, c, root, "")
-
- return list(res)
+ for r in range(rows):
+ for c in range(columns):
+ dfs(r, c, trie.root, "")
+ return list(result)