Leetcode 212. Word Search II. Python (TrieNode/Prefix Tree)

 212. Word Search II

Given an m x n board of characters and a list of strings words, return all words on the board.

Each word must be constructed from letters of sequentially adjacent cells, where adjacent cells are horizontally or vertically neighboring. The same letter cell may not be used more than once in a word.

Input: board = [["o","a","a","n"],["e","t","a","e"],["i","h","k","r"],["i","f","l","v"]], words = ["oath","pea","eat","rain"]
Output: ["eat","oath"]

Constraints:

• m == board.length
• n == board[i].length
• 1 <= m, n <= 12
• board[i][j] is a lowercase English letter.
• 1 <= words.length <= 3 * 104
• 1 <= words[i].length <= 10
• words[i] consists of lowercase English letters.
• All the strings of words are unique.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.eow = False
        self.counter = 0
    
    def wadd(self, w):
        curr = self
        curr.counter += 1
        for c in w:
            if c not in curr.children:
                curr.children[c] = TrieNode()
            curr = curr.children[c]
            curr.counter += 1
        curr.eow = True
    
    def wremove(self, w):
        curr= self
        curr.counter -= 1
        for c in w:
            if c in curr.children:
                curr = curr.children[c]
                curr.counter -= 1
    
class Solution:
    def findWords(self, b: List[List[str]], words: List[str]) -> List[str]:
        root = TrieNode()
        for w in words:
            root.wadd(w)
        
        rows, cols = len(b), len(b[0])
        ans, vs = set(), set()
        
        def dfs(r,c,n,w):
            if (r<0 or c<0 or 
                r==rows or c==cols or
               (r,c) in vs or b[r][c] not in n.children or
                n.children[b[r][c]].counter < 1):
                return
            
            vs.add((r,c))
            n = n.children[b[r][c]]
            w += b[r][c]
            if n.eow:
                n.eow = False
                ans.add(w)
                root.wremove(w)
                
            dfs(r+1,c,n,w)
            dfs(r-1,c,n,w)
            dfs(r,c+1,n,w)
            dfs(r,c-1,n,w)
            
            vs.remove((r,c))
            
        for r in range(rows):
            for c in range(cols):
                dfs(r,c,root,"")
                
        return list(ans)
            

Explaination :