word square generator

wordsquare.py

from random import sample
from itertools import product

def make_trie(words):
    root = dict()
    for word in words:
        parent = root
        for letter in word:
            parent = parent.setdefault(letter, dict())
    return root

def make_word_square(size=5):
    with open('/usr/share/dict/words') as fh:
        words = [w.strip() for w in fh]
        
    root = make_trie(w.lower() for w in words if len(w) == size)
        
    down_trie = [[None] * size for _ in range(size)]
    across_trie = [[None] * size for _ in range(size)]
    letters = [[''] * size for _ in range(size)]
    
    for i in range(size):
        down_trie[0][i] = root
        across_trie[i][0] = root
    
    coords = list(sorted(product(range(size), range(size)), key=sum))
    def fill_square(i):
        row, col = coords[i]
        #print('position: {} {}'.format(row, col))
        #for rr in letters:
        #    print(' '.join(rr))
        dt = down_trie[row][col]
        at = across_trie[row][col]
        choices = list(set(at) & set(dt))

        shuffle(choices)
        
        if not choices:
            letters[row][col] = ''
            return False
        
        for choice in choices:
            letters[row][col] = choice

            if row < size - 1:
                down_trie[row+1][col] = dt[choice]
            if col < size - 1:
                across_trie[row][col+1] = at[choice]
        
            if i == len(coords) - 1:
                return True
            else:
                if fill_square(i+1):
                    return True

        letters[row][col] = ''
        return False
    
    fill_square(0)

    for rr in letters:
        print(' '.join(rr))