Treap


Treap is a binary search tree that combines the properties of a binary search tree and a heap. It maintains the binary search tree property with respect to the keys and the heap property with respect to the priorities assigned to the keys. The key priority is randomly assigned during the insertion process and the treap is balanced by performing rotations based on the priorities.

Let's implement a basic Treap data structure in Python:



import random

class Node:
    def __init__(self, key):
        self.key = key
        self.priority = random.random()
        self.left = None
        self.right = None

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

    def rotate_left(self, node):
        y = node.right
        node.right = y.left
        y.left = node
        return y

    def rotate_right(self, node):
        x = node.left
        node.left = x.right
        x.right = node
        return x

    def insert(self, key):
        if self.root is None:
            self.root = Node(key)
            return

        self.root = self._insert(self.root, key)

    def _insert(self, node, key):
        if node is None:
            return Node(key)

        if key < node.key:
            node.left = self._insert(node.left, key)
            if node.left.priority < node.priority:
                node = self.rotate_right(node)
        else:
            node.right = self._insert(node.right, key)
            if node.right.priority < node.priority:
                node = self.rotate_left(node)

        return node

    def inorder_traversal(self, node):
        if node is None:
            return

        self.inorder_traversal(node.left)
        print(node.key, end=' ')
        self.inorder_traversal(node.right)

# Create a Treap
treap = Treap()

# Insert nodes into the Treap
treap.insert(5)
treap.insert(3)
treap.insert(7)
treap.insert(1)
treap.insert(4)

# Perform inorder traversal to print the Treap
print("Treap (Inorder Traversal):")
treap.inorder_traversal(treap.root)


In this example, we define a Node class to represent the nodes of the Treap and a Treap class to implement the Treap operations. The insert method inserts a new node into the Treap based on the key value and assigns a random priority value. If the priority violates the heap property, rotations are performed to restore the heap property. The rotate_left and rotate_right methods perform left and right rotations, respectively. The inorder_traversal method performs an inorder traversal of the Treap to print its contents. We create a Treap, insert nodes into it, and print the Treap using inorder traversal. 

Now let's solve a problem using the Treap data structure. The problem we'll solve is "Count of Range Sum", which asks to count the total number of range sums that fall within a given range.

class TreapNode:
    def __init__(self, key, priority):
        self.key = key
        self.priority = priority
        self.count = 1
        self.left = None
        self.right = None

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

    def rotate_left(self, node):
        # Perform a left rotation on the given node
        y = node.right
        node.right = y.left
        y.left = node
        return y

    def rotate_right(self, node):
        # Perform a right rotation on the given node
        x = node.left
        node.left = x.right
        x.right = node
        return x

    def insert(self, key, priority):
        self.root = self._insert(self.root, key, priority)

    def _insert(self, node, key, priority):
        if node is None:
            return TreapNode(key, priority)

        if key == node.key:
            node.count += 1
        elif key < node.key:
            node.left = self._insert(node.left, key, priority)
            if node.left.priority > node.priority:
                node = self.rotate_right(node)
        else:
            node.right = self._insert(node.right, key, priority)
            if node.right.priority > node.priority:
                node = self.rotate_left(node)

        return node

    def count_range_sum(self, nums, lower, upper):
        prefix_sums = [0] * (len(nums) + 1)
        for i in range(len(nums)):
            prefix_sums[i + 1] = prefix_sums[i] + nums[i]

        count = 0
        sum_tree = Treap()
        for i in range(len(prefix_sums)):
            count += sum_tree.count_range(prefix_sums[i] - upper, prefix_sums[i] - lower)
            sum_tree.insert(prefix_sums[i], random.random())

        return count

    def count_range(self, lower, upper):
        return self._count_range(self.root, lower, upper)

    def _count_range(self, node, lower, upper):
        if node is None:
            return 0

        if lower <= node.key <= upper:
            return node.count + self._count_range(node.left, lower, upper) + self._count_range(node.right, lower, upper)
        elif node.key < lower:
            return self._count_range(node.right, lower, upper)
        else:
            return self._count_range(node.left, lower, upper)
 

The time complexity of the insert and delete operations in the Treap data structure is O(log n) on average, and the search operation is also O(log n) on average. The space complexity is O(n) to store the Treap nodes.

Next Post Previous Post