Find the number of elements larger than a given element in BST

Question

I'm trying to solve the Hackerrank's Insertion Sort Advanced Analysis problem using BST (similar to this question on SO). As I put items in the tree, I need to find out the number of items greater than it (the get_rank() method in the code):

class Node:
    def __init__(self, data):
        self.left = None
        self.right = None
        self.data = data
        self.num_left_children = 0
        self.num_right_children = 0

    def insert(self, data):
        if data <= self.data:
            if self.left is None:
                self.left = Node(data)
            else:
                self.left.insert(data)
            self.num_left_children += 1
        else:
            if self.right is None:
                self.right = Node(data)
            else:
                self.right.insert(data)
            self.num_right_children += 1

    def get_rank(self, data):
        if data < self.data:
            return self.num_right_children + self.left.get_rank(data) + 1
        elif data > self.data:
            return self.right.get_rank(data)
        else:
            return self.num_right_children

How can I improve the performance of this code (e.g. in case multiple identical items are put into the tree) ?

Answer 1

Answering my own question, I made these improvements to the code which allowed me to solve the HR problem:

• Instead of representing duplicate values as separate nodes use a counter of occurrences of the value.

• Since we need to get the rank of a value immediately after inserting it, we can combine the insert() and get_rank() methods.

• Turn recursion into iteration.

The final code:

class Node:
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None
        self.num_right_children = 0
        self.occurrences = 1

    def insert(self, data):
        current_node = self
        rank = 0

        while True:
            if data < current_node.data:
                rank += current_node.num_right_children + current_node.occurrences
                if current_node.left is None:
                    current_node.left = Node(data)
                    break
                current_node = current_node.left
            elif data > current_node.data:
                current_node.num_right_children += 1
                if current_node.right is None:
                    current_node.right = Node(data)
                    break
                current_node = current_node.right
            else:
                current_node.occurrences += 1
                rank += current_node.num_right_children
                break
        return rank