Coding for Interviews: Outward Counterclockwise Spiral Matrix Traversal

spiral.py

UP    = (-1,  0)
LEFT  = ( 0, -1)
DOWN  = ( 1,  0)
RIGHT = ( 0,  1)

class Turtle:
    """The turtle crawls over the matrix starting at some position.
    The turtle knows how to turn left and go forward."""
    def __init__(self, nrows, ncols, x, y):
        self.x = x
        self.y = y
        self.dy, self.dx = UP
        self.nrows = nrows
        self.ncols = ncols
        self.traversal = []

    def turn_left(self):
        """Turn left. Does not change the current location of the turtle."""
        self.dy, self.dx = {UP: LEFT, LEFT: DOWN, DOWN: RIGHT, RIGHT: UP}[(self.dy, self.dx)]

    def go_forward(self):
        """Move forward in the current direction.
        This can move the turtle off the matrix."""
        self.x += self.dx
        self.y += self.dy
        if 0 <= self.y < self.nrows and 0 <= self.x < self.ncols:
            self.traversal.append(self.nrows*self.y + self.x + 1)

    def seen_everything(self):
        """Returns True if the turtle has traversed the entire matrix."""
        return len(self.traversal) == self.nrows*self.ncols

def spiral_traversal(nrows, ncols, row, col):
    """Perform a left spiral traversal of a (nrows,ncols) matrix starting at (row,col).
    Co-ordinates are provided using 1-based indexing.
    Returns a list of the traversed elements."""
    #
    # The example uses 1-based indexing for matrix co-ordinates.
    # Internally, we use 0-based indexing for everything, so handle the discrepancy here.
    #
    t = Turtle(nrows, ncols, col-1, row-1)
    edge_counter = 0
    edge_length = 1
    while not t.seen_everything():
        for step in range(edge_length):
            t.go_forward()
        t.turn_left()
        #
        # The edge lengths go like this:
        # 1 1 2 2 3 3 4 4 ...
        #
        edge_counter += 1
        if edge_counter % 2 == 0:
            edge_length += 1
    return t.traversal
            
def main():
    import sys
    a, b, c, d = map(int, sys.argv[1:])
    print spiral_traversal(a, b, c, d)
    
if __name__ == "__main__":
    main()