bst.py (plain text)


"""The bst module contains functions implementing the immutable binary search
tree (BST) data abstraction for CS61A at UC Berkeley.

BSTs are (immutable) binary search trees, a common example of a tree structure
that is used for storing sorted data in a way to allow for quick searching
later.

One might notice that none of the doctests for either ADT actually "show" an
BST. This is because this ADT is meant to be treated as a black box: the
tests do not assume any specific implementation for BST.
"""


# General Binary Search Trees Abstraction


empty_bst = None


def make_bst(datum, left=empty_bst, right=empty_bst):
    """Create a BST with datum, left branch, and right branch. 

    It should be noted that nothing in the constructor or selectors actually
    force the BST to be a proper search tree, so in reality this makes just a
    simple binary tree.

    >>> b1 = make_bst(0, left=make_bst(-1), right=make_bst(1))
    >>> b2 = make_bst(0, right=make_bst(1))
    >>> b3 = make_bst(0, left=make_bst(-1))
    """
    return (left, datum, right)


def bst_datum(b):
    """Return the datum of BST b.

    >>> b1 = make_bst(0, left=make_bst(-1), right=make_bst(1))
    >>> bst_datum(b1)
    0
    """
    return b[1]


def bst_left(b):
    """Return the left branch of BST b.

    >>> b1 = make_bst(0, left=make_bst(-1), right=make_bst(1))
    >>> bst_datum(bst_left(b1))
    -1
    """
    return b[0]


def bst_right(b):
    """Return the right branch of BST b.

    >>> b1 = make_bst(0, left=make_bst(-1), right=make_bst(1))
    >>> bst_datum(bst_right(b1))
    1
    """
    return b[2]


# General Binary Search Trees Utility Functions


def bst_is_leaf(b):
    """Return True if BST b is a leaf.

    >>> bst_is_leaf(make_bst(5))
    True
    >>> bst_is_leaf(make_bst(5, left=make_bst(4)))
    False
    """
    return bst_left(b) is empty_bst and bst_right(b) is empty_bst


def bst_find(b, item):
    """Return True if item is in BST b and False otherwise.

    >>> test_bst = make_bst(0,
    ...                     left=make_bst(-4,
    ...                                   left=make_bst(-6,
    ...                                                 left=make_bst(-7),
    ...                                                 right=make_bst(-5)),
    ...                                   right=make_bst(-2,
    ...                                                  left=make_bst(-3),
    ...                                                  right=make_bst(-1))),
    ...                     right=make_bst(4,
    ...                                    left=make_bst(2,
    ...                                                  left=make_bst(1),
    ...                                                  right=make_bst(3)),
    ...                                    right=make_bst(6,
    ...                                                   left=make_bst(5),
    ...                                                   right=make_bst(7))))
    >>> bst_find(test_bst, 5)
    True
    >>> bst_find(test_bst, 22)
    False
    >>> bst_find(test_bst, 0.5)
    False
    >>> bst_find(test_bst, 1)
    True
    """
    if b == empty_bst:
        return False
    elif bst_datum(b) == item:
        return True
    elif bst_datum(b) > item:
        return bst_find(bst_left(b), item)
    return bst_find(bst_right(b), item)


def bst_insert(b, item):
    """Insert an item into the BST b (if it's not already in the tree).

    >>> test_bst = make_bst(3,
    ...                     left=make_bst(1,
    ...                                   left=make_bst(0)),
    ...                     right=make_bst(5,
    ...                                    right=make_bst(6)))
    >>> print(bst_str(test_bst))
    3
    |
    +-1
    | |
    | +-0
    | |
    | +-empty_bst
    |
    +-5
      |
      +-empty_bst
      |
      +-6
    >>> print(bst_str(bst_insert(test_bst, 2)))
    3
    |
    +-1
    | |
    | +-0
    | |
    | +-2
    |
    +-5
      |
      +-empty_bst
      |
      +-6
    >>> print(bst_str(bst_insert(test_bst, 7)))
    3
    |
    +-1
    | |
    | +-0
    | |
    | +-empty_bst
    |
    +-5
      |
      +-empty_bst
      |
      +-6
        |
        +-empty_bst
        |
        +-7
    """
    if b == empty_bst:
        return make_bst(item)
    elif bst_datum(b) == item:
        return b # Don't insert duplicates
    elif bst_datum(b) > item:
        return make_bst(bst_datum(b),
                        bst_insert(bst_left(b), item),
                        bst_right(b))
    return make_bst(bst_datum(b),
                    bst_left(b),
                    bst_insert(bst_right(b), item))

def bst_str(b, depth=0, left_child=False):
    """Return a string representation of the BST b

    Uses empty_bst to represent empty branches if only one Branch is empty.
    Admittedly, not the prettiest or most intuitive way to represent BSTs, but
    we use this to simplify the code.

    >>> test_bst = make_bst(3,
    ...                     left=make_bst(1,
    ...                                   left=make_bst(0)),
    ...                     right=make_bst(5,
    ...                                    right=make_bst(6)))
    >>> print(bst_str(test_bst))
    3
    |
    +-1
    | |
    | +-0
    | |
    | +-empty_bst
    |
    +-5
      |
      +-empty_bst
      |
      +-6
    """
    result = repr(bst_datum(b))
    if not bst_is_leaf(b):
        # Left Child
        result += "\n|\n+-"
        if bst_left(b) is not empty_bst:
            result += bst_str(bst_left(b), depth + 1, left_child=True)
        else:
            result += "empty_bst"
        # Right Child
        result += "\n|\n+-"
        if bst_right(b) is not empty_bst:
            result += bst_str(bst_right(b), depth + 1, left_child=False)
        else:
            result += "empty_bst"

    # Add front padding if we're part of a bigger tree we're converting to a
    # string.  Perhaps not the most elegant code.
    if depth > 0:
        if left_child:
            leader = "| "
        else:
            leader = "  "
        result = result.replace("\n", "\n" + leader)
    return result