Tower of Hanoi

hanoi.py

import sys
import unittest
import copy


class HanoiTest(unittest.TestCase):
    def test_number_of_steps_3(self):
        ''' number of steps should be 2^n-1 (3 disks) '''
        h = Hanoi(3).solve()
        self.assertEqual(len(h), 2**3)

    def test_number_of_steps_7(self):
        ''' number of steps should be 2^n-1 (7 disks) '''
        h = Hanoi(7).solve()
        self.assertEqual(len(h), 2**7)

    def test_number_of_steps_10(self):
        ''' number of steps should be 2^n-1 (10 disks) '''
        h = Hanoi(10).solve()
        self.assertEqual(len(h), 2**10)

    def test_trivial_case(self):
        ''' 1 disk should go directly from A to C '''
        h = Hanoi(1).solve()
        self.assertEqual(h[0]['state'], {'A':[1], 'B':[], 'C':[]})
        self.assertEqual(h[1]['state'], {'A':[], 'B':[], 'C':[1]})
    
    @unittest.skip('TODO')
    def test_disk_order(self):
        ''' disks should always be in sorted order '''
        pass


class Hanoi:
    def __init__(self, height):
        self.height = height
        self.stack = {}
        self.steps = []

    def _move_to(self, height, from_stack, via_stack, to_stack):
        if height >= 1:
            self._move_to(height-1, from_stack, to_stack, via_stack)
            self._move_disk(from_stack, to_stack)
            self._move_to(height-1, via_stack, from_stack, to_stack)

    def _move_disk(self, from_stack, to_stack):
        self.stack[to_stack].append(self.stack[from_stack].pop())
        self.steps.append({'move': (from_stack, to_stack), 'state': copy.deepcopy(self.stack)})

    def _pad(self, height, stack):
            return reversed(['='] + [str(i) for i in stack] + [' '] * (height - len(stack)))

    def print_steps(self):
        for step in self.steps:
            stack_a = self._pad(self.height, step['state']['A'])
            stack_b = self._pad(self.height, step['state']['B'])
            stack_c = self._pad(self.height, step['state']['C'])
            for st in zip(stack_a, stack_b, stack_c):
                print '{0}   {1}   {2}'.format(st[0], st[1], st[2])
            print '{0}   {1}   {2}'.format('A', 'B', 'C')
            print '  {0} to {1}'.format(step['move'][0], step['move'][1])

    def solve(self):
        self.stack = {'A': range(self.height, 0, -1), 'B': [], 'C': []}
        self.steps = [{'move': ('None', 'None'), 'state': copy.deepcopy(self.stack)}]
        self._move_to(self.height, 'A', 'B', 'C')
        return self.steps


if __name__ == '__main__':
    height = int(sys.argv[1])
    hanoi = Hanoi(height)
    hanoi.solve()
    hanoi.print_steps()