Solution to Calvin traffic light puzzle

Calvin_lights.py

#! /usr/bin/env python

# Calvin has to cross several signals when he walks from his home to school.
# Each of these signals operate independently.
# They alternate every 80 seconds between green light and red light.
# At each signal, there is a counter display that tells him how long it will be before the current signal light changes.
# Calvin has a magic wand which lets him turn a signal from red to green instantaneously.
# However, this wand comes with limited battery life, so he can use it only for a specified number of times.
#
# a. If the total number of signals is 2 and Calvin can use his magic wand only once, then what is the expected waiting
# time at the signals when Calvin optimally walks from his home to school?
#
# b. What if the number of signals is 3 and Calvin can use his magic wand only once?
#
# c. Can you write a code that takes as inputs the number of signals and the number of times Calvin can use his magic
# wand, and outputs the expected waiting time?

import argparse

from anytree import Node, RenderTree
from blessings import Terminal

t = Terminal()


class CalvinNode(Node):
    def __init__(self, name, parent=None, levels=2, probability=1.0, period=80, total=0.0, wand=1):
        super(CalvinNode, self).__init__(name, parent=parent)
        self.levels = levels
        self.left = None
        self.right = None
        # Problem variables
        self.probability = probability
        self.period = period
        self.total = total
        self.wand = wand

    def _get_path(self):
        return "{}".format(self.separator.join([""] + [str(node.name) for node in self.path]))

    def __repr__(self):
        return "Node('{}', levels={}, probability={}, total={})".format(
            self._get_path(), self.levels, self.probability, self.total,
        )


class TrafficNode(CalvinNode):
    def __init__(self, name, parent=None,
                 levels=2, probability=1.0, period=80, total=0.0, wand=1):
        super(TrafficNode, self).__init__(
            name, parent=parent,
            levels=levels, probability=probability, period=period, total=total, wand=wand
        )
        if levels == 0:
            return
        else:
            self.right = TrafficNode(t.green('G'), parent=self,
                                     levels=self.levels - 1, probability=self.probability / 2,
                                     period=self.period,
                                     total=0.0,
                                     wand=self.wand)
        if self.wand == 0:
            if levels >= 1:
                self.left = TrafficNode(t.red('R'), parent=self,
                                        levels=self.levels - 1, probability=self.probability / 2,
                                        period=self.period,
                                        total=(self.period / 2) * (self.probability / 2),
                                        wand=self.wand)
        elif self.wand > 0:
            if self.wand < self.levels:
                self.left = WandNode(t.red('R'), parent=self,
                                     levels=self.levels, probability=self.probability / 2,
                                     period=self.period,
                                     wand=self.wand)
            else:
                self.left = TrafficNode(t.red('R') + '-' + t.yellow('used'), parent=self,
                                        levels=self.levels - 1, probability=self.probability / 2,
                                        period=self.period,
                                        wand=self.wand - 1)
        self.total = self.total + self.left.total + self.right.total

    def __repr__(self):
        return "traffic" + super(TrafficNode, self).__repr__()


class WandNode(CalvinNode):
    def __init__(self, name, parent=None,
                 levels=2, probability=1.0, period=80, total=0.0, wand=1):
        super(WandNode, self).__init__(
            name, parent=parent,
            levels=levels, probability=probability, period=period, total=total, wand=wand
        )
        self.left = TrafficNode(t.cyan('unused'), parent=self,
                                levels=self.levels - 1, probability=self.probability / 4,
                                period=self.period,
                                wand=self.wand)
        # Since we use the wand when threshold exceeds 20s, the cost of not using it averages 10s
        self.left.total = (self.left.period / 2) / 4 * self.left.probability
        self.right = TrafficNode(t.yellow('used'), parent=self,
                                 levels=self.levels - 1, probability=self.probability / 4 * 3,
                                 period=self.period,
                                 wand=self.wand - 1)
        self.total = self.left.total + self.right.total

    def __repr__(self):
        return "wand" + super(WandNode, self).__repr__()


def test_calvin():
    calvin = TrafficNode('calvin',
                         levels=2,
                         wand=1)
    assert calvin.total == 8.75


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-l", "--levels", type=int, required=True,
                        help="Number of traffic lights")
    parser.add_argument("-w", "--wand", type=int, required=True,
                        help="Number of wand uses")
    args = parser.parse_args()

    calvin = TrafficNode('calvin_{}_{}'.format(args.levels, args.wand),
                         levels=args.levels,
                         wand=args.wand)
    print
    print RenderTree(calvin)
    print