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

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

Author: mitchellirvin

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)