Python program for Word Search II

This is a Leetcode problem –

Given a 2D board and a list of words from the dictionary, find all words in the board.

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

Note:

  • All inputs are consist of lowercase letters a-z.
  • The values of words are distinct.

Here is my solution to this challenge –

class Solution:     def __init__(self, board, words):         self.board = board         self.words = words      def find_words(self, board, words):          root = {}         for word in words:             node = root             for c in word:                 node = node.setdefault(c, {})             node[None] = True         board = {i + 1j * j: c                  for i, row in enumerate(board)                  for j, c in enumerate(row)}          found = []         def search(node, z, word):             if node.pop(None, None):                 found.append(word)             c = board.get(z)             if c in node:                 board[z] = None                 for k in range(4):                     search(node[c], z + 1j ** k, word + c)                 board[z] = c         for z in board:             search(root, z, '')          return found 

Program explanation – I first build a tree of words with root root and also represent the board a different way, namely as a one-dimensional dictionary where the keys are complex numbers representing the row/column indexes. That makes further work with it easier. Looping over all board positions is just for z in board, the four neighbors of a board position z are just z + 1j ** k (for k in 0 to 3), and I don’t need to check borders because board.get just returns None if I request an invalid position.

After this preparation, I just take the tree and recursively dive with it into each board position. Similar to how you’d search a single word, but with the tree instead.

Here is an example input/output –

output = Solution([   ['o','a','a','n'],   ['e','t','a','e'],   ['i','h','k','r'],   ['i','f','l','v'] ], ["oath","pea","eat","rain"])  print(output.find_words([   ['o','a','a','n'],   ['e','t','a','e'],   ['i','h','k','r'],   ['i','f','l','v'] ], ["oath","pea","eat","rain"]))  >>> ['oath', 'eat'] 

So, I would like to know whether I could make this program shorter and more efficient.

Any help would be highly appreciated.