test_tree.py

import sys
import unittest
import logging

from tree import Tree, tree, max_depth, max_depth_lr

logging.basicConfig(level=logging.DEBUG)


class TestInitTree(unittest.TestCase):
    def test_init_empty(self):
        t = tree(0)
        assert isinstance(t, Tree)
        self.assertEqual(t.x, 0)
        self.assertEqual(t.l, None)
        self.assertEqual(t.r, None)

    def test_init_two_node(self):
        t = tree(0, tree(1), None)
        assert isinstance(t, Tree)
        self.assertEqual(t.x, 0)
        assert isinstance(t.l, Tree)
        self.assertEqual(t.l.x, 1)
        self.assertEqual(t.r, None)


class TestMaxDepth(unittest.TestCase):
    def assertDepth(self, depth, t):
        found = max_depth(t)
        self.assertEqual(depth, found,
                "Expected %d found %d for %s" % (
                    depth, found, t))

    def test_single(self):
        self.assertDepth(0, tree())

    def test_two_node(self):
        logging.debug("one")
        self.assertDepth(1, tree(0, tree()))
        logging.debug("two")
        self.assertDepth(1, tree(0, None, tree()))

    def test_three_node_two_level(self):
        self.assertDepth(1, tree(0, tree(), tree()))

    def test_three_level_left(self):
        self.assertDepth(2, tree(0, tree(0, tree(), None), None))

    def test_four_level_complex(self):
        t = tree(0,
                tree(1, tree(2), tree(2)),
                tree(1, tree(2,
                    tree(3, None, None),
                    None),
                    )
                )
        self.assertDepth(3, t)


class TestMaxDepthLR(unittest.TestCase):
    def assertDepth(self, depth, t):
        found = max_depth_lr(t)
        self.assertEqual(depth, found,
                "Expected %d found %d for %s" % (
                    depth, found, t))

    def test_single(self):
        self.assertDepth(0, tree())

    def test_two_level(self):
        self.assertDepth(1, tree(0, tree()))
        self.assertDepth(1, tree(0, None, tree()))

    def test_three_level_l(self):
        t = tree(0, tree(1, tree(2)))
        self.assertDepth(2, t)

    def test_three_level_r(self):
        t = tree(0, None, tree(1, None, tree(2)))
        self.assertDepth(2, t)

    def test_three_level_switch(self):
        t = tree(0, tree(1, None, tree(2)))
        self.assertDepth(1, t)


if __name__ == '__main__':
    unittest.main()

tree.py

import logging


class Tree(object):
    x = 0
    l = None
    r = None

    def __str__(self):
        return "%r" % self

    def __repr__(self):
        return "tree(%d, %r, %r)" % (
                self.x, self.l, self.r)


def tree(*args):
    t = Tree()
    if len(args) > 0:
        t.x = args[0]
    if len(args) > 1:
        t.l = args[1]
    if len(args) > 2:
        t.r = args[2]
    logging.debug("tree(%r) returning %s" % (
        args, t))
    return t

def solution(T):
    return max_depth(T)

def max_depth(T):
    """
    Use a depth-first traversal to determine the maximum depth of tree
    T.
    """

    if T.l is None:
        max_depth_l = 0
    else:
        max_depth_l = 1 + max_depth(T.l)

    if T.r is None:
        max_depth_r = 0
    else:
        max_depth_r = 1 + max_depth(T.r)

    return max(max_depth_r, max_depth_l)

def max_depth_lr(T, direction=None):
    if direction == 'L' and T.l is not None:
        depth = 1 + max_depth_lr(T.l, 'L')
    elif direction == 'R' and T.r is not None:
        depth = 1 + max_depth_lr(T.r, 'R')
    elif direction == None:
        depth_l = 1 + max_depth_lr(T.l, 'L') if T.l is not None else 0
        depth_r = 1 + max_depth_lr(T.r, 'R') if T.r is not None else 0
        depth = max(depth_l, depth_r)
    else:
        depth = 0
    return depth