Computing Rank of a node in a stream

rank_node_2.py

# Computing Rank of a node in a stream

"""Source: https://www.geeksforgeeks.org/rank-element-stream """

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

# Inserting a new Node
def insert(root, val):
    if root is None:
        return Node(val)
     
    # updating size of left subtree
    if val < root.val:
        root.left = insert(root.left, val)
        root.left_size += 1
    else:
        root.right = insert(root.right, val)
    
    return root

def get_rank(root, x):
    """Get rank of Node with value X """

    if root.val == x:
        return root.left_size
    
    elif x < root.val:
        if root.left is None:
            return -1
        else:
            return get_rank(root.left, x)
        
    else:
        if root.right is None:
            return -1
        else:
            right_size = get_rank(root.right, x)
            if right_size == -1:
                return -1
            else:
                return root.left_size + (1+ right_size)
    

if __name__=='__main__':
    arr = [5, 1, 4, 4, 5, 9, 7, 13, 3]
    n = len(arr)
    x = 4

    root = None
    for i in range(n):
        root = insert(root, arr[i])
    
    print("Rank of ", x, "in stream is:", get_rank(root, x))

    x = 13
    print("Rank of ", x, "in stream is:", get_rank(root, x))

    x = 8
    print("Rank of ", x, "in stream is:", get_rank(root, x))