I tried to implement the binary tree program in python. I want to add one more function to get the level of a specific node.

eg:-

         10         #level 0
         /  \
       5    15     #level 1
      /  \
     3   7         #level 2

if we search the level of node 3, it should return

class BinaryTree:
    def __init__(self, data):
        self.left = None
        self.right = None
        self.data = data
        
    def insert(self, data):
        if self.data == data:
            return
        if self.data<data:
            if self.left:
                self.left.insert(data)
            else:
                self.left = BinaryTree(data)
        else:
            if self.right:
                self.right.insert(data)
            else:
                self.right = BinaryTree(data)
        
    def print_tree(self):
        if self:
            print(self.data)    
        if self.left:
            self.left.print_tree()
        elif self.right:
            self.right.print_tree()
            
    def get_level(self, data, level=0):
        print(type(data))
        level  = 1
        if self.data == data:
            return level
        elif self.left and self.right:
            if self.left:
                self.left.get_level(data)
            elif self.right:
                self.right.get_level(data)
        return level
    
    def in_order(self):
        if self:
            #left
            in_order(self.left)
            #root
            print(self.data, '->')
            #right
            in_order(self.right)

This is my code for Binary tree. I need to add another function to this class which will tell the level of a node after giving the value

For example:

def get_level(node):
    #some code
    return level


get_level(10) # returns the level of the node 10

CodePudding user response：

This worked for me:

def get_level(self, data, level=0):
    print(type(data))
    level  = 1
    if self.data == data:
        return level
    if data > self.data:
        if self.left:
            return self.left.get_level(data,level)
    else:
        if self.right:
            return self.right.get_level(data,level)

    return None

For the 10 the level will be 1, for 5 and 15 it is level 2, and for 3 and 7 it is level 3.

for the 0,1,2 levels just remove the "level =1" line and add a 1 in level variable at the recursion call.

CodePudding user response：

Before getting into your question, there are some issues in your existing code:

• The insert method is using the BST property of the tree (good), but is going to the wrong subtree to find the location of the insertion point. The condition if self.data<data should be inverted to if self.data > data.

• The function inorder has recursive calls that call the function not as a method but as a global function. So that means the indentation of this function is wrong. It should be defined outside of the class, not inside.

I need to add another function to this class

You already have that function, but it has these issues:

• The function returns a level, but when you make the recursive call, you ignore that returned value.
• As you never pass the level argument, that level will always start with 0 and so the returned value will always be 1.
• With the elif self.left and self.right you ignore the case where a node has one child.
• No benefit is taken from the fact that this is a BST, and so both left and right subtrees are traversed, while data can be found by following a single path.

I would suggest to do this without level parameter, and let the function return the level as if the given node were the root. The caller can than add 1 after the recursive call.

    def get_level(self, data):
        if data == self.data:
            return 0
        level = None
        if data < self.data:
            if self.left:
                level = self.left.get_level(data)
        else:
            if self.right:
                level = self.right.get_level(data)
        if level is not None:
            return level   1