CODEWITHSUNDEEP
pythonbinary-search-tree
oneman
oneman

Reputation: 811

How to implement 'range' with a python BST

My code at the moment allows me to search for a specific node, I would like to edit it so that I can search within a range of numbers. For example, I have a price list of apples, and I would like to add all apples to a list/dictionary that cost between $2-$4 or something like that.

Here is my current code

def valueOf(self, key):
    node = self._bstSearch(self._root, key)
    assert node is not None, "Invalid map key."
    return node.value

def _bstSearch(self, subtree, target):
    if subtree is None:
        return None
    elif target < subtree.key:
        return self._bstSearch( subtree.left, target)
    elif target > subtree.key:
        return self._bstSearch(subtree.right, target)
    else:    
        return subtree

I think I should be editing target to change it from a single item search to a ranged item search but I'm not 100% sure how to

Upvotes: 1

Views: 1207

Answers (1)

Mikhail M.
Mikhail M.

Reputation: 5988

Using recursive in-order traversal:

def range(self, a, b):
    return self._traverse_range(self._root, a, b)

def _traverse_range(self, subtree, a, b, cumresult=None):
    if subtree is None:
        return

    # Cumulative variable.
    if cumresult is None:
        cumresult = []

    # Traverse LEFT subtree if it is possible to find values in required range there.
    if subtree.key > a:
        self._traverse_range(subtree.left, a, b, cumresult)

    # Push VALUE if it is in our range.
    if a <= subtree.key <= b:  # Change to strict "< b" to act like python's range
        cumresult.append(subtree.key)

    # Traverse RIGHT subtree if it is possible to find values in required range there.
    if subtree.key < b:
        self._traverse_range(subtree.right, a, b, cumresult)

    return cumresult

Upvotes: 3

Related Questions