Balanced Trees

Master Treap, split/merge operations, and self-balancing tree techniques

Overview

Balanced trees are efficient data structures for maintaining ordered data collections. This chapter explores Treap insertion/deletion, split/merge operations, and various self-balancing tree techniques for dynamic ordered set maintenance.

Core Content

1. Basic Treap Operations

// Treap implementation
#include <iostream>
#include <random>
using namespace std;

struct TreapNode {
    int key, priority, size;
    TreapNode* left;
    TreapNode* right;
    
    TreapNode(int k) : key(k), priority(rand()), size(1), left(nullptr), right(nullptr) {}
};

class Treap {
private:
    TreapNode* root;
    
    int getSize(TreapNode* node) {
        return node ? node->size : 0;
    }
    
    void updateSize(TreapNode* node) {
        if (node) {
            node->size = 1 + getSize(node->left) + getSize(node->right);
        }
    }
    
    // Right rotation
    TreapNode* rotateRight(TreapNode* node) {
        TreapNode* newRoot = node->left;
        node->left = newRoot->right;
        newRoot->right = node;
        updateSize(node);
        updateSize(newRoot);
        return newRoot;
    }
    
    // Left rotation
    TreapNode* rotateLeft(TreapNode* node) {
        TreapNode* newRoot = node->right;
        node->right = newRoot->left;
        newRoot->left = node;
        updateSize(node);
        updateSize(newRoot);
        return newRoot;
    }
    
    TreapNode* insert(TreapNode* node, int key) {
        if (!node) return new TreapNode(key);
        
        if (key < node->key) {
            node->left = insert(node->left, key);
            if (node->left->priority > node->priority) {
                node = rotateRight(node);
            }
        } else if (key > node->key) {
            node->right = insert(node->right, key);
            if (node->right->priority > node->priority) {
                node = rotateLeft(node);
            }
        }
        
        updateSize(node);
        return node;
    }
    
    TreapNode* remove(TreapNode* node, int key) {
        if (!node) return nullptr;
        
        if (key < node->key) {
            node->left = remove(node->left, key);
        } else if (key > node->key) {
            node->right = remove(node->right, key);
        } else {
            if (!node->left) {
                TreapNode* temp = node->right;
                delete node;
                return temp;
            } else if (!node->right) {
                TreapNode* temp = node->left;
                delete node;
                return temp;
            } else {
                if (node->left->priority > node->right->priority) {
                    node = rotateRight(node);
                    node->right = remove(node->right, key);
                } else {
                    node = rotateLeft(node);
                    node->left = remove(node->left, key);
                }
            }
        }
        
        updateSize(node);
        return node;
    }
    
    // Find k-th smallest
    int kthSmallest(TreapNode* node, int k) {
        if (!node) return -1;
        
        int leftSize = getSize(node->left);
        if (k == leftSize + 1) {
            return node->key;
        } else if (k <= leftSize) {
            return kthSmallest(node->left, k);
        } else {
            return kthSmallest(node->right, k - leftSize - 1);
        }
    }
    
public:
    Treap() : root(nullptr) {}
    
    void insert(int key) {
        root = insert(root, key);
    }
    
    void remove(int key) {
        root = remove(root, key);
    }
    
    int kthSmallest(int k) {
        return kthSmallest(root, k);
    }
    
    int size() {
        return getSize(root);
    }
};

int main() {
    Treap treap;
    
    // Insert elements
    treap.insert(5);
    treap.insert(3);
    treap.insert(8);
    treap.insert(1);
    treap.insert(7);
    
    cout << "3rd smallest element: " << treap.kthSmallest(3) << endl;
    cout << "Tree size: " << treap.size() << endl;
    
    treap.remove(3);
    cout << "After removing 3, 3rd smallest: " << treap.kthSmallest(3) << endl;
    
    return 0;
}

2. Split and Merge Operations

// Treap split/merge implementation
class SplitMergeTreap {
private:
    struct Node {
        int key, priority, size;
        Node* left;
        Node* right;
        
        Node(int k) : key(k), priority(rand()), size(1), left(nullptr), right(nullptr) {}
    };
    
    Node* root;
    
    void updateSize(Node* node) {
        if (node) {
            node->size = 1 + getSize(node->left) + getSize(node->right);
        }
    }
    
    int getSize(Node* node) {
        return node ? node->size : 0;
    }
    
    // Split by value
    pair<Node*, Node*> split(Node* node, int key) {
        if (!node) return {nullptr, nullptr};
        
        if (node->key <= key) {
            auto [left, right] = split(node->right, key);
            node->right = left;
            updateSize(node);
            return {node, right};
        } else {
            auto [left, right] = split(node->left, key);
            node->left = right;
            updateSize(node);
            return {left, node};
        }
    }
    
    // Split by size
    pair<Node*, Node*> splitBySize(Node* node, int k) {
        if (!node) return {nullptr, nullptr};
        
        int leftSize = getSize(node->left);
        if (leftSize >= k) {
            auto [left, right] = splitBySize(node->left, k);
            node->left = right;
            updateSize(node);
            return {left, node};
        } else {
            auto [left, right] = splitBySize(node->right, k - leftSize - 1);
            node->right = left;
            updateSize(node);
            return {node, right};
        }
    }
    
    // Merge two treaps
    Node* merge(Node* left, Node* right) {
        if (!left) return right;
        if (!right) return left;
        
        if (left->priority > right->priority) {
            left->right = merge(left->right, right);
            updateSize(left);
            return left;
        } else {
            right->left = merge(left, right->left);
            updateSize(right);
            return right;
        }
    }
    
public:
    SplitMergeTreap() : root(nullptr) {}
    
    void insert(int key) {
        auto [left, right] = split(root, key);
        Node* newNode = new Node(key);
        root = merge(merge(left, newNode), right);
    }
    
    void remove(int key) {
        auto [left, temp] = split(root, key - 1);
        auto [mid, right] = split(temp, key);
        root = merge(left, right);
    }
    
    // Range operation: get range [l,r]
    vector<int> getRange(int l, int r) {
        auto [left, temp] = splitBySize(root, l - 1);
        auto [mid, right] = splitBySize(temp, r - l + 1);
        
        vector<int> result;
        inorder(mid, result);
        
        root = merge(merge(left, mid), right);
        return result;
    }
    
private:
    void inorder(Node* node, vector<int>& result) {
        if (!node) return;
        inorder(node->left, result);
        result.push_back(node->key);
        inorder(node->right, result);
    }
};

int main() {
    SplitMergeTreap treap;
    
    for (int i = 1; i <= 10; i++) {
        treap.insert(i);
    }
    
    vector<int> range = treap.getRange(3, 7);
    cout << "Elements in range [3,7]: ";
    for (int x : range) {
        cout << x << " ";
    }
    cout << endl;
    
    return 0;
}

Problem-Solving Tips

🎯 Operation Selection

Choose appropriate operations based on problem requirements: use basic operations for insertion/deletion, split/merge for range operations.

⚡ Performance Optimization

Maintain subtree size information for k-th smallest queries, use random priorities to ensure expected performance.

🔧 Implementation Points

Pay attention to updating node information, handle null pointers correctly, maintain BST properties in split/merge operations.

🧮 Application Scenarios

Dynamically maintain ordered sets, support insertion/deletion, k-th smallest queries, range operations, and other complex requirements.