Binary Tree in Python Number of Nodes in Range

I have a simple binary tree implementation:

class Node:
    def __init__(self, item, left = None, right = None):
        self.item = item
        self.left = left
        self.right = right

class BST:
    def __init__(self):
        self.root = None

    def add(self, item):
        if self.root == None:
            self.root = Node(item, None, None)
        else:
            child_tree = self.root
            while child_tree != None:
                parent = child_tree
                if item < child_tree.item:
                    child_tree = child_tree.left
                else:
                    child_tree = child_tree.right
            if item < parent.item:
                parent.left = Node(item, None, None)
            elif item > parent.item:
                parent.right = Node(item, None, None)

I want to add the count (lo, hi) method, which counts all the nodes in the range (lo, hi) (including hi). This is what I still have:

def count(self, lo, hi, ptr='lol', count=0):
    if ptr == 'lol':
        ptr = self.root
    if ptr.left != None:
        if ptr.item >= lo and ptr.item <= hi:
            count += 1
        ptr.left = self.count(lo, hi, ptr.left, count)
    if ptr.right != None:
        if ptr.item >= lo and ptr.item <= hi:
            count += 1
        ptr.right = self.count(lo, hi, ptr.right, count)
    return count

It seems that this only works when the binary tree is right lean or left lean. This does not work for balanced trees, and I have no idea why. My input:

bst = BST()
for ele in [10, 150, 80, 40, 20, 10, 30, 60, 50, 70, 120, 100, 90, 110, 140, 130, 150]:
    bst.add(ele)
print(bst.count(30, 100))

My code gives me output: 0, but it has to say output: 8. Can you tell me where I was wrong, please?

+4
source share
1 answer

Wrong part:

   while child_tree != None:
        if child_tree.item >= lo and child_tree.item <= hi:
            count += 1
        if hi > child_tree.item:  # from here
            child_tree = child_tree.right
        else:
            child_tree = child_tree.left . # to here

child_tree hi, - .

: , ...

UPDATE

def count(self, lo, hi, ptr, count=0):
    if not ptr:
        return 0
    elif lo <= ptr.item <= hi:
        return 1 + self.count(lo, hi, ptr.left, count) + \
               self.count(lo, hi, ptr.right, count)
    elif ptr.item < lo:
        return self.count(lo, hi, ptr.right, count)
    elif ptr.item > hi:
        return self.count(lo, hi, ptr.left, count)
+2

Source: https://habr.com/ru/post/1687984/


All Articles