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.