binary trees in python

binary_search_tree.py

class Node(object):
  """A node referencing other nodes, forming a subtree."""

  def __init__(self, data):
    self.left = None
    self.right = None
    self.data = data

  def __repr__(self):
    """A (not-so-good) display."""
    result = ''
    if self.left:
      result += str(self.left)
    result += '%s ' % str(self.data)
    if self.right:
      result += str(self.right)
    return result

  def insert(self, data):
    """Add a node."""
    if not self.data:
      self.data = data
      return
    if data < self.data:
      if not self.left:
        self.left = Node(data)
      else:
        self.left.insert(data)
    elif data > self.data:
      if not self.right:
        self.right = Node(data)
      else:
        self.right.insert(data)

  def lookup(self, data):
    """Find a node."""
    if data < self.data:
      if not self.left:
        return None
      return self.left.lookup(data)
    elif data > self.data:
      if not self.right:
        return None
      return self.right.lookup(data)
    return self.data

  def get_parent(self, target):
    """Returns (parent_data, side) where side is 'left' or 'right'."""
    # Check self.
    if self.data == target:
      return None
    # Check children.
    if self.left and self.left.data == target:
      return self.data
    if self.right and self.right.data == target:
      return self.data
    # No dice, search the left and right sides.
    if self.left:
      result = self.left.get_parent(target)
      if result:
        return result
    if self.right:
      result = self.right.get_parent(target)
      if result:
        return result
    # Nothing doing!
    return None

  def get_children(self, target):
    """Returns a tuple representing the (left, right) child Nodes."""
    # Check self.
    if self.data == target:
      result = []
      if self.left:
        result.append(self.left.data)
      else:
        result.append(None)
      if self.right:
        result.append(self.right.data)
      else:
        result.append(None)
      return tuple(result)
    # Search left.
    if self.left:
      result = self.left.get_children(target)
      if result:
        return result
    # Search right.
    if self.right:
      result = self.right.get_children(target)
      if result:
        return result
    return None

  def count_nodes(self):
    """Count the number of nodes in the tree, including self."""
    count = 1
    if self.left:
      count += self.left.count_nodes()
    if self.right:
      count += self.right.count_nodes()
    return count

  def delete(self, target):
    """Delete a Node and its children."""
    # Can't delete the root..
    if self.data == target:
      if not self.get_parent(self.data):
        raise ValueError
    # Search the children.
    if self.left and self.left.data == target:
      self.left = None
    if self.right and self.right.data == target:
      self.right = None
    # Search left and right trees.
    if self.left:
      self.left.delete(target)
    if self.right:
      self.right.delete(target)

  def equals(self, tree):
    """Returns True if self and 'tree' are equal, False otherwise."""
    # This node DNE or has different data.
    if not tree or not tree.data or tree.data != self.data:
      return False
    # The branches have different numbers of nodes.
    if self.count_nodes() != tree.count_nodes():
      return False
    # Compare left for existence and data-sameness.
    if self.left and not tree.left:
      return False
    if tree.left and not self.left:
      return False
    if self.left:
      return self.left.equals(tree.left)
    # Compare right.
    if self.right and not tree.right:
      return False
    if tree.right and not self.right:
      return False
    if self.right:
      return self.right.equals(tree.right)
    # No differences found.
    return True

test_binary_search_tree.py

import unittest

import binary_search_tree as tree


class TreeTest(unittest.TestCase):

  def build_tree(self):
    root = tree.Node(8)
    root.insert(3)
    root.insert(10)
    root.insert(1)
    root.insert(6)
    root.insert(4)
    root.insert(7)
    root.insert(14)
    root.insert(13)
    return root

  def setUp(self):
    self.root = self.build_tree()

  def test_lookup(self):
    """We can lookup Nodes by their data."""
    self.assertEqual(10, self.root.lookup(10))

  def test_get_parent(self):
    """We can find parent Nodes, if they exist."""
    self.assertEqual(3, self.root.get_parent(6))
    self.assertEqual(6, self.root.get_parent(4))
    self.assertEqual(8, self.root.get_parent(10))
    self.assertEqual(10, self.root.get_parent(14))
    self.assertEqual(None, self.root.get_parent(8))
    self.assertEqual(None, self.root.get_parent(500))

  def test_get_children(self):
    """We can get data on Node children."""
    self.assertEqual((1, 6), self.root.get_children(3))
    self.assertEqual((None, 14), self.root.get_children(10))
    self.assertEqual((13, None), self.root.get_children(14))
    self.assertEqual((None, None), self.root.get_children(7))

  def test_count_nodes(self):
    """We can count the number of nodes."""
    self.assertEqual(9, self.root.count_nodes())

  def test_delete(self):
    """Nodes (and their associated sub-branches) can be dropped."""
    self.root.delete(7)
    self.assertEqual((4, None), self.root.get_children(6))
    self.root.delete(10)
    self.assertEqual(None, self.root.lookup(14))
    with self.assertRaises(ValueError):
      self.root.delete(8)

  def test_equals(self):
    """Two trees can be compared."""
    second_tree = self.build_tree()
    self.assertTrue(self.root.equals(second_tree))
    second_tree.delete(10)
    self.assertFalse(self.root.equals(second_tree))
    # Make two new trees with the same Node count but different values.
    tree_one = tree.Node(5)
    tree_one.insert(3)
    tree_one.insert(10)
    tree_one.insert(1)
    tree_two = tree.Node(5)
    tree_two.insert(3)
    tree_two.insert(10)
    tree_two.insert(2)
    self.assertFalse(tree_one.equals(tree_two))