Word Squares Generator

Question

I've solved the Word Squares problem:

Given a set of words (without duplicates), find all word squares you can build from them.

A sequence of words forms a valid word square if the kth row and column read the exact same string, where 0 ≤ k < max(numRows, numColumns).

For example, the word sequence ["ball","area","lead","lady"] forms a word square because each word reads the same both horizontally and vertically.

 b a l l
 a r e a
 l e a d
 l a d y

but the following code, had a runtime so slow that I ended up in the bottom 10% of submissions speed wise. I was wondering if there's anything I'm doing in the code that's a performance killer / any way I could improve my algorithm to speed it up.

from collections import defaultdict

class Solution(object):
    def make_defaultdict(self):
        return defaultdict(self.make_defaultdict)

    def make_trie(self, words):
        trie = self.make_defaultdict()

        for word in words:
            cur = trie
            for l in word:
                cur = cur[l]
            cur['end'] = True

        return trie

    def search_trie(self, trie, search_str):
        matches = []
        sub_trie = trie

        for l in search_str:
            if l not in sub_trie:
                return []
            sub_trie = sub_trie[l]

        q = [(search_str + k, sub_trie[k]) for k in sub_trie.keys()]
        while q:
            fragment, sub_trie = q.pop(0)

            if 'end' in sub_trie:
                matches.append(fragment)
            else:
                q.extend([(fragment + k, sub_trie[k]) for k in sub_trie.keys()])

        return matches

    def make_square(self, pos, square, size, trie):
        if pos == size:
            return [square]

        for i in range(pos, size):
            search_str = ''.join([w[i] for w in square])
            candidates = self.search_trie(trie, search_str)

            if not candidates:
                return []

            result = []
            for candidate in candidates:
                result.extend(self.make_square(pos + 1, square + [candidate], size, trie))

            return result

    def wordSquares(self, words):
        all_squares = []
        trie = self.make_trie(words)

        for word in words:
            all_squares.extend(self.make_square(1, [word], len(word), trie) )

        return all_squares

Any thoughts on how to speed it up?

Answer 1

1. Review

1. There are no docstrings. What do any of these functions do? How do I call them? What do they return?

2. Just because leetcode.com insists on a Solution class, doesn't mean that all functions need to be members of that class. Write plain functions and avoid the unnecessary use of self.

3. I like the implementation of make_defaultdict, but this doesn't just make a defaultdict, it makes an empty prefix tree (trie), so a name like make_empty_prefix_tree would be better. (I prefer the name "prefix tree" because it's clear what it is from the name, whereas the name "trie" is a bit mysterious.)

4. Distinguishing the end of the prefix_tree by a special key like 'end' ought not to be necessary in this application (constructing word squares) because all of the words are the same length.

5. When you need to iterate over the keys and values of dictionary simultaneously, use dict.items. For example, this line:

   q = [(search_str + k, sub_trie[k]) for k in sub_trie.keys()]
   

   becomes

   q = [(search_str + k, v) for k, v in sub_trie.items()]
   

   avoiding the lookup of sub_trie[k].

2. Performance

In order to understand what the code is doing, it is handy to trace its execution on an example. This can be done using the Python debugger, for example we might print the value of square on each call to make_square, like this:

>>> pdb.run("Solution().wordSquares('tide idea dear yarn tidy'.split())")
> <string>(1)<module>()
(Pdb) break Solution.make_square
Breakpoint 1 at cr176612.py:38
(Pdb) commands 1
(com) silent
(com) p ' '.join(square)
(com) continue
(Pdb) continue
'tide'
'tide idea'
'tide idea dear'
'idea'
'idea dear'
'dear'
'yarn'
'tidy'
'tidy idea'
'tidy idea dear'
'tidy idea dear yarn'

(You could get a similar result by editing the code to add a print at the start of make_square but it is a good idea to practice using the debugger on easy examples like this, so that when you have a tough bug to investigate you are proficient with it.)

Looking at the output carefully you can see the problem: after choosing tide for the first row, the search continued by choosing idea for the second row. But this should have been ruled out because then the fourth column would have to start ea and nothing in the dictionary will fit there.

So work is wasted following search paths that could have been rejected earlier if all constraints had been applied. The problem is that only one constraint is applied at a time — on the \$i\$th loop in make_square we use search_trie to apply the constraint for the \$i\$th column only, failing to apply the constraints from columns \$j > i\$.

3. Organizing the work

It takes some thought to figure out how to organize the work so that all constraints are applied at the earliest possible stage. This is the kind of algorithm where it makes sense to work out the details on paper for a simple example, before starting to write code.

Suppose our dictionary consists of all, led, and lee, making the following prefix tree:

The algorithm is going to need the idea of "current position in the prefix tree". I'm going to draw this with a red triangle. For example, if l and e have been matched, then the current position looks like this:

In this position, d and e are the possible matches for the next letter.

Now, in a 3×3 word square there are nine squares to be filled:

In the real problem the result is required to be symmetric, so that choice 2 is the same as choice 4 and so on, but it will be clearer, I think, to solve the general problem first and then figure out how to take advantage of symmetry later.

When we come to make a choice for a square there are two constraints we have to satisfy, a horizontal constraint for the row and a vertical constraint for the column. The constraints correspond to positions in the prefix tree. For choice number 1, the constraints in both directions correspond to the root of the prefix tree, like this:

The vertical constraint says that we can choose a or l, and so does the horizontal constraint. Let's suppose that we choose a. Then we can update the position in the prefix tree to make constraints for the next cells, like this:

Now we can consider the constraints bearing on choice number 2, bringing in a fresh vertical constraint:

Here, the vertical constraint says that we can choose a or l, and the horizontal constraint says that we can only choose l, so this is our only way to satisfy both constraints. Once again we can propagate the constraints:

The process continues through the choices in order, until all constraints have been applied and a solution has been found (or until a dead end has been reached where the constraints cannot be satisfied, in which case we have to backtrack). If a solution is found, the data structures will look something like this:

4. Data structures

Given the approach outlined in §3 above, how can we represent the state of the search in a data structure? Well, we will need choices (the black squares in the diagram below) which we will consider in order. For each choice we have constraints bearing on that choice (the blue diamonds connected to the choice by blue lines), and there are rules specifying how to propagate the constraints after a choice has been made (the red arrows which join two blue diamonds by passing through a black square).

We can represent constraints by objects with a node attribute:

class Constraint:
    node = None     # Node in the prefix tree, or None if not initialized.

(These need to be objects rather than, say, named tuples, because the node gets modified when a constraint is propagated.)

We can represent choices as objects with lists of constraints and propagation rules:

class Choice:
    def __init__(self, constraints, rules):
        """Create a choice.

        Arguments:
        constraints: List of Constraint objects.
        rules: List of propagation rules as pairs (Constraint, Constraint).

        """
        self.constraints = constraints
        self.rules = rules

    def options(self):
        "Return set of available options, assuming initialized constraints."
        # Intersect all the keys of the constraint dictionaries.
        constraints = iter(self.constraints)
        options = next(constraints).node.keys()
        for c in constraints:
            options &= c.node.keys()
        return options

    def propagate(self, choice):
        "Propagate a choice from source to destination constraints."
        for source, dest in self.rules:
            dest.node = source.node[choice]

Now finding all assignments that meet the constraints is a straightforward depth-first search, and a convenient way to implement this in Python is to use the stack of iterators pattern, like this:

def assignments(choices):
    """Generate satisfying assignments to a sequence of choices."""
    n = len(choices)
    assignment = [None]
    stack = [iter(choices[0].options())]

    # Cache methods in local variables to avoid lookups in the loop
    assignment_append = assignment.append
    assignment_pop = assignment.pop
    stack_append = stack.append
    stack_pop = stack.pop

    while stack:
        for option in stack[-1]:
            i = len(assignment)
            if i < n:
                assignment_append(option)
                choices[i - 1].propagate(option)
                stack_append(iter(choices[i].options()))
                break
            yield tuple(assignment[1:]) + (option,)
        else:
            assignment_pop()
            stack_pop()

Notice that this generates the assignments one by one instead of collecting them all and returning a list. Generating results is more flexible than returning a list: for example, if you only want one assignment then you can use next(assignments(choices)), or if you want to see whether there is a unique assignment or not then you can use islice(assignments(choices), 2), without having to wait for the search to run all the way to completion.

It's a bit fiddly to build the data structure, but we can do it like this:

from itertools import product

def word_squares(words):
    """Generate word squares from a sequence of words."""
    n = len(words[0])
    assert all(len(word) == n for word in words)
    prefix_tree = make_prefix_tree(words)

    # Horizontal and vertical constraints for each choice.
    N = range(n)
    horiz, vert = [[[Constraint() for _ in N] for _ in N] for _ in range(2)]
    for k in N:
        horiz[k][0].node = vert[0][k].node = prefix_tree

    # Create choices.
    choices = []
    for i, j in product(range(n), repeat=2):
        h, v = horiz[i][j], vert[i][j]
        rules = []
        if j < n - 1:
            rules.append((h, horiz[i][j + 1]))
        if i < n - 1:
            rules.append((v, vert[i + 1][j]))
        choices.append(Choice([h, v], rules))

    # Find assignments and reformat them.
    for assignment in assignments(choices):
        yield [''.join(word) for word in zip(*[iter(assignment)] * n)]

5. Adding symmetry

Now it ought to be clear how to implement symmetry. We drop the squares that we don't need to fill, and reroute the constraint propagation like this:

It's still fiddly to build the choices structure, but it can be done like this:

from itertools import combinations_with_replacement

def symmetric_word_squares(words):
    """Generate symmetric word squares from a sequence of words."""
    n = len(words[0])
    assert all(len(word) == n for word in words)
    prefix_tree = make_prefix_tree(words)

    # Horizontal and vertical constraints for each choice (some we
    # won't use, but it's simpler like this).
    N = range(n)
    horiz, vert = [[[Constraint() for _ in N] for _ in N] for _ in range(2)]
    for k in N:
        horiz[k][0].node = vert[0][k].node = prefix_tree

    # Create choices.
    choices = []
    for i, j in combinations_with_replacement(N, 2):
        h, v = horiz[i][j], vert[i][j]
        rules = []
        if i == j:
            constraints = [v]
            if j < n - 1:
                rules.append((v, horiz[i][j + 1]))
        else:
            constraints = [h, v]
            if i != j < n - 1:
                rules.append((h, horiz[i][j + 1]))
            if j != i < n - 1:
                rules.append((v, vert[i + 1][j]))
        choices.append(Choice(constraints, rules))

    # Find assignments.
    grid = [[None] * n for _ in N]
    for assignment in assignments(choices):
        for a, (i, j) in zip(assignment, combinations_with_replacement(N, 2)):
            grid[i][j] = grid[j][i] = a
        yield [''.join(word) for word in grid]

Note the use of itertools.combinations_with_replacement here — this is useful when you would otherwise need two nested loops to loop over all \$0 \le i \le j < n\$:

    for i in range(N):
        for j in range(i, N):

Here combinations_with_replacement combines the two loops into one, saving a level of indentation.

6. Timing comparisons

I got kind of carried away there with data structures and algorithms, so I'm not sure if the revised code will really be any faster. Let's try an example, using all the 4-letter words in my system dictionary, which can make 4.3 million symmetric word squares:

>>> words = [w for w in map(str.strip, open('/usr/share/dict/words'))
...          if len(w) == 4 and w == w.lower()]
>>> len(words)
4347
>>> sum(1 for _ in symmetric_word_squares(words))
4264868
>>> timeit(lambda:sum(1 for _ in symmetric_word_squares(words)), number=1)
49.30328282299888
>>> timeit(lambda:Solution().wordSquares(words), number=1)
66.61988531300449

So the revised code takes about 75% as long as the original.