Program to find out the node in the right in a binary tree using Python

Sometimes we need to find the node immediately to the right of a given node in a binary tree. The right node must be at the same level as the target node. This problem can be solved using level-order traversal (BFS) to process nodes level by level.

So, if the input is like

and u = 6, then the output will be 8. The node situated at the right of node 6 is node 8, so the value 8 is returned.

Algorithm

To solve this, we will follow these steps −

• if root is empty, return null

• Use a deque for level-order traversal

• For each level, store all nodes in a temporary list

• Find the target node's position in the current level

• Return the next node if it exists, otherwise return null

Implementation

from collections import deque

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def insert(temp, data):
    que = []
    que.append(temp)
    while len(que):
        temp = que[0]
        que.pop(0)
        if not temp.left:
            if data is not None:
                temp.left = TreeNode(data)
            else:
                temp.left = TreeNode(0)
            break
        else:
            que.append(temp.left)
        if not temp.right:
            if data is not None:
                temp.right = TreeNode(data)
            else:
                temp.right = TreeNode(0)
                break
        else:
            que.append(temp.right)

def make_tree(elements):
    tree = TreeNode(elements[0])
    for element in elements[1:]:
        insert(tree, element)
    return tree

def search_node(root, element):
    if root is None:
        return None
    if root.val == element:
        return root
    res1 = search_node(root.left, element)
    if res1:
        return res1
    res2 = search_node(root.right, element)
    return res2

def find_right_node(root, u):
    if not root:
        return None
    
    dq = deque()
    dq.append(root)
    
    while dq:
        dq_size = len(dq)
        temp = []
        index = -1
        
        for _ in range(dq_size):
            node = dq.popleft()
            if node.left:
                dq.append(node.left)
            if node.right:
                dq.append(node.right)
            temp.append(node)
            if node == u:
                index = len(temp) - 1
        
        # If target node is the last node in this level, no right node exists
        if index == len(temp) - 1:
            return None
        # If target node found, return the next node
        if index > -1:
            return temp[index + 1]
    
    return None

# Create the binary tree
root = make_tree([5, 3, 7, 2, 4, 6, 8])
u = search_node(root, 6)
result = find_right_node(root, u)

if result:
    print(result.val)
else:
    print("No right node found")

8

How It Works

The algorithm uses level-order traversal (BFS) to process nodes level by level. For each level, it stores all nodes in a temporary list and tracks the position of the target node. If the target node is found and it's not the last node in that level, the algorithm returns the next node in the list.

Example with Different Target

# Same tree, but looking for node to the right of 4
root = make_tree([5, 3, 7, 2, 4, 6, 8])
u = search_node(root, 4)
result = find_right_node(root, u)

if result:
    print(f"Right node of 4: {result.val}")
else:
    print("No right node found for 4")

Right node of 4: 6

Conclusion

This solution efficiently finds the right neighbor of any node in a binary tree using level-order traversal. The time complexity is O(n) where n is the number of nodes, and space complexity is O(w) where w is the maximum width of the tree.