Python solution to Cracking the Coding Interview Question 17.25 - Word Rectangle

WordRectangle.py

### WordRectangle.py
### Python solution to Cracking the Coding Interview's question 17.25
### Noah Nelson
###
### The problem:
### Given a dictionary or list of words, return the largest rectangle you can
### form out of those words where every row and every column is a valid word.
### All of the rows have to have the same length, and all of the columns,
### but the width and height may be different. The largest rectangle is the
### one with the largest area.

import sys
from math import ceil, sqrt


##############################
### Prefix Tree classes 
###

class TreeNode:
    """
    A node in a prefix tree. Has a letter, as well as a dictionary of child
    letters mapped to the treenodes they represent.
    """

    def __init__(self, letter, endsWord=False):
        self.letter = letter
        self.children = {}
        # endsWord indicates if this node is the end of a word, not just
        # an orphaned prefix.
        self.endsWord = endsWord

    def addChild(self, child):
        """
        Adds the given child to this treenode, if there isn't already a child
        tree for its letter.
        """
        if child.letter not in self.children:
            self.children[child.letter] = child

    def addWord(self, word):
        """
        Adds a word to this tree node by recursively adding the tail of the
        word to the appropriate child node.
        """
        if len(word) == 0:
            self.endsWord = True
            return
        head = word[0]
        tail = word[1:]
        self.addChild(TreeNode(head))
        self.children[head].addWord(tail)

    def searchPrefix(self, prefix):
        """
        Searches the given word recursively to see if the entire prefix is in
        the tree rooted at this node. Expects the prefix as some type
        that supports slicing, list or string.
        """
        if len(prefix) == 0:
            return True
        head = prefix[0]
        tail = prefix[1:]
        return head in self.children and\
                self.children[head].searchPrefix(tail)

    def searchWord(self, word):
        """
        Searches the given word recursively to see if the word is in the tree
        rooted at this node. Only differs from searchPrefix in the recursive
        base case - empty word is valid only if this node represents the end
        of the word. An example: If your tree contained just "dog", then "do"
        is a valid prefix but not a valid word.
        """
        if len(word) == 0:
            return self.endsWord
        head = word[0]
        tail = word[1:]
        return head in self.children and self.children[head].searchWord(tail)


class PrefixTree:
    """
    A dictionary of valid words represented as a prefix tree.

    Checks if a string is a prefix of any valid word in O(k) time where
    k is the length of the prefix - bounded by a constant for English and
    probably any reasonable language, even German.
    """

    def __init__(self, words=[]):
        """Create a prefix tree containing all the words in the iterator."""
        self.root = TreeNode(None) # Special root treenode with no letter
        for word in words:
            self.addWord(word)

    def addWord(self, word):
        """Adds a given word and all of its prefixes to the tree."""
        self.root.addWord(word)
    
    def isPrefix(self, prefix):
        """
        Check if a given prefix is a valid prefix to any words in the
        dictionary this prefix tree represents.
        """
        return self.root.searchPrefix(prefix)

    def isWord(self, word):
        """
        Check if a given word is a valid word.
        This is different from isPrefix, for example because "kn" is the
        prefix to the word "know" but is not a word itself, depending on your
        words file.
        """
        return self.root.searchWord(word)

##############################
### Rectangle class
###
### Used to hide implementation details from rectangle searching functions,
### because I experimented with tons of different implementations of this
### class.
###

class Rectangle:

    def __init__(self, length):
        """Create an empty rectangle of the given length."""
        self.words = []
        self.length = length
        self.width = 0

    def addWord(self, word):
        """Add the given word as the next row in this rectangle."""
        assert(len(word) == self.length)
        self.words.append(word)
        self.width += 1

    def removeWord(self):
        """Remove the last row of the rectangle."""
        self.width -= 1
        self.words.pop()

    def columns(self):
        """Returns a generator of the columns of the rectangle."""
        return ([word[j] for word in self.words] for j in range(self.length))

    def prettyPrint(self):
        """Print this prettily."""
        for word in self.words:
            print(word)



def _validRectangle(predicate):
    """
    Returns a function to check a rectangle's validity according to a given
    predicate. This is used to create the partial and full rectangle validity
    functions with respect to some prefix tree.

    Only checks if the rectangle's columns all satisfy the predicate, since
    rows are always assumed to be valid words.
    """

    def valid(rect):
        # Assume valid rows and just form all column words
        return all((predicate(col) for col in rect.columns()))

    return valid


##############################
### makeRectangle function
###
### The main important function - recrusive DFS through possible rectangles.
###

def makeRectangle(width, fullPred, partialPred, words, rect):
    """
    Attempt to make a rectangle with the given length and width, using
    a given list of words of the rectangle's length.

    Recursive, so we pass in a partially complete rectangle object.

    Arguments:
    width: target width of the rectangle.
    fullPred: predicate which returns true iff the rectangle is fully valid
    partialPred: predicate which tells if the partial rectangle is valid so far
    words: a sequence of words that are the same length as the rectangle
    rect: a partially complete rectangle object to build from

    Returns:
    a rectangle object with the given width and the length of the rect passed
    in, or None if none exists using the given words.
    """
    if rect.width == width:
        if fullPred(rect):
            return rect
        else:
            return None

    for word in words:
        # Try adding each word to the end, then recursing on the result.
        rect.addWord(word)
        if not partialPred(rect):
            # Adding that word created an invalid partial column. Try next word
            rect.removeWord()
            continue
        else:
            thatRect = makeRectangle(width, fullPred, partialPred, words, rect)
            if thatRect:
                return thatRect
            rect.removeWord()
            # Remove this failed word so we can try the next word.

    # If all those fail, it's impossible.
    return None

def largestRectangle(words):
    """
    Get the largest rectangle that can be formed with the given list of
    words. Returns either a rectangle object, or None if none could be
    formed with those words.

    This function groups the words by length, then walks down the possible
    areas we could achieve, checking if each one has a rectangle with that
    area.
    """

    # Process the words to split them up by length.
    # Build the prefix tree at the same time so we only go through the
    # words generator once.
    tree = PrefixTree()
    lengthTable = {}
    for word in words:
        tree.addWord(word)

        length = len(word)
        if length in lengthTable:
            lengthTable[length].append(word)
        else:
            lengthTable[length] = [word]

    longestWord = max(lengthTable.keys())

    # Predicates for if a rectangle is partially or fully valid are with
    # reference to the prefix tree we created.
    fullRect = _validRectangle(tree.isWord)
    partialRect = _validRectangle(tree.isPrefix)

    # Note: an improvement could be made to reduce redundant searches.
    # When we attempt to form a 10x10 rectangle, we're also trying 10x8
    # and 10x9 and all the way down sizes. So we could get the information
    # about whether there's a 10x9 rectangle when we search the tree of
    # possible 10x10 rectangles.

    # Now, walk backwards, attempting to form a rectangle of each area.
    targetArea = longestWord * longestWord
    while targetArea > 0:
        for l in reversed(range(ceil(sqrt(targetArea)), longestWord+1)):
            if targetArea % l == 0 and targetArea / l <= longestWord:
                #print("Looking for a {l} x {w} rectangle...".format(
                #        l=l, w=targetArea // l))
                words = lengthTable[l] if l in lengthTable else []
                rect = makeRectangle(targetArea // l,
                            fullRect, partialRect, words, Rectangle(l))
                if rect:
                    return rect
        targetArea -= 1

    # If we haven't returned, we found nothing.
    return None


##############################
### Main entry point.
###

usageString = """
Usage: {} <words file | -t>

Attempt to form the largest possible crossword rectangle with the words in a
given file, or just run some simple tests if the -t argument is supplied.
""".format(sys.argv[0])

if __name__ == '__main__':

    if len(sys.argv) != 2:
        print(usageString)
        exit(1)

    if sys.argv[1] == "-t":
        print("Running the simple tests of the PrefixTree object.")
        words = ['noah', 'nelson', 'noble', 'james', 'tim', 'time', 'jonny',
                 'ian', 'jonni']
        t = PrefixTree(words)
        for prefix in ['no', 'yes', 'ti', 'jonn', 'jonnh', 'x', 'jam']:
            print(t.isPrefix(prefix))
        print("expected T, F, T, T, F, F, T")
        for word in ['noah', 'noble', 'nobility', 'brogogog', 'x', 'jon']:
            print(t.isWord(word))
        print("exprected T, T, F, F, F, F")
        exit(0)

    filename = sys.argv[1]

    rect = largestRectangle(word[:-1] for word in open(filename, 'r'))
    if rect:
        print("found a rectangle of area {}".format(rect.length * rect.width))
        rect.prettyPrint()
    else:
        print("No valid rectangle found. :(")