Word Search II

Problem: Word Search II

We use a Trie to solve this problem. We first check every element in board to see any word in Trie start with that character, then iterate through its neighbors and search down in Trie to see anything matching.

Code in Python:

class Solution(object):
    def findWords(self, board, words):
        """
        :type board: List[List[str]]
        :type words: List[str]
        :rtype: List[str]
        """
        trie = Trie()
        for word in words:
            trie.insert(word)

        ans, visited = [], []

        def checkTrie(node, s, cur, i, j):
            if node.children[ord(cur)-ord('a')]:
                if node.children[ord(cur)-ord('a')].isWord:
                    if s + cur not in ans:
                        ans.append(s + cur)
                visited.append((i, j))
                if (i-1, j) not in visited and i > 0:
                    checkTrie(node.children[ord(cur)-ord('a')], s+cur, board[i-1][j], i-1, j)
                if (i+1, j) not in visited and i < len(board)-1:
                    checkTrie(node.children[ord(cur)-ord('a')], s+cur, board[i+1][j], i+1, j)
                if (i, j-1) not in visited and j > 0:
                    checkTrie(node.children[ord(cur)-ord('a')], s+cur, board[i][j-1], i, j-1)
                if (i, j+1) not in visited and j < len(board[0])-1:
                    checkTrie(node.children[ord(cur)-ord('a')], s+cur, board[i][j+1], i, j+1)
                visited.pop()
            else:
                return

        for i in xrange(len(board)):
            for j in xrange(len(board[0])):
                checkTrie(trie.root, "", board[i][j], i, j)
        return ans


class TrieNode(object):
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.isWord = False
        self.children = [None] * 26


class Trie(object):

    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        """
        Inserts a word into the trie.
        :type word: str
        :rtype: void
        """
        cur = self.root
        for c in word:
            if cur.children[ord(c)-ord('a')] == None:
                cur.children[ord(c)-ord('a')] = TrieNode()
            cur = cur.children[ord(c)-ord('a')]
        cur.isWord = True