AVLTree.hpp

#pragma once

#include <iostream>
#include <vector>
#include <string>
#include <algorithm> // For std::max
#include <memory>    // For std::unique_ptr

/*
 A node in AVL Tree
 K Key type (e.g., post timestamp)
 V Value type (e.g., post content string)
*/
template <typename K, typename V>
struct AVLNode {
    K key;
    V value;
    std::unique_ptr<AVLNode<K, V>> left;
    std::unique_ptr<AVLNode<K, V>> right;
    int height;

    AVLNode(K k, V v) : key(k), value(v), left(nullptr), right(nullptr), height(1) {}
};

/*
 AVL Tree Implementation,, Stores posts sorted by creation time.
 K Key type (timestamp)
 V Value type (post content)
*/
template <typename K, typename V>
class AVLTree {
private:
    std::unique_ptr<AVLNode<K, V>> root;
    int getHeight(const std::unique_ptr<AVLNode<K, V>>& node) {
        return node ? node->height : 0;
    }
    int getBalance(const std::unique_ptr<AVLNode<K, V>>& node) {
        return node ? (getHeight(node->left) - getHeight(node->right)) : 0;
    }
    void updateHeight(std::unique_ptr<AVLNode<K, V>>& node) {
        if (node) {
            node->height = 1 + std::max(getHeight(node->left), getHeight(node->right));
        }
    }
    //Right Rotation
    std::unique_ptr<AVLNode<K, V>> rightRotate(std::unique_ptr<AVLNode<K, V>> y) {
        std::unique_ptr<AVLNode<K, V>> x = std::move(y->left);
        std::unique_ptr<AVLNode<K, V>> T2 = std::move(x->right);

        x->right = std::move(y);
        x->right->left = std::move(T2); // y->left = T2

        updateHeight(x->right);
        updateHeight(x);
        return x;
    }
    //Left Rotation
    std::unique_ptr<AVLNode<K, V>> leftRotate(std::unique_ptr<AVLNode<K, V>> x) {
        std::unique_ptr<AVLNode<K, V>> y = std::move(x->right);
        std::unique_ptr<AVLNode<K, V>> T2 = std::move(y->left);

        y->left = std::move(x);
        y->left->right = std::move(T2); // x->right = T2

        updateHeight(y->left);
        updateHeight(y);
        return y;
    }
    //Inserting Key,Value pair using Recursion
    std::unique_ptr<AVLNode<K, V>> insert(std::unique_ptr<AVLNode<K, V>> node, K key, V value) {
        if (!node) {
            return std::make_unique<AVLNode<K, V>>(key, value);
        }
        if (key < node->key) {
            node->left = insert(std::move(node->left), key, value);
        } else if (key > node->key) {
            node->right = insert(std::move(node->right), key, value);
        } else {
            // No duplication of keys for timestamps
            return node;
        }

        updateHeight(node);
        int balance = getBalance(node);

        //Rebalancing (if needed)
        // Left Left Case
        if (balance > 1 && key < node->left->key) {
            return rightRotate(std::move(node));
        }
        // Right Right Case
        if (balance < -1 && key > node->right->key) {
            return leftRotate(std::move(node));
        }
        // Left Right Case
        if (balance > 1 && key > node->left->key) {
            node->left = leftRotate(std::move(node->left));
            return rightRotate(std::move(node));
        }
        // Right Left Case
        if (balance < -1 && key < node->right->key) {
            node->right = rightRotate(std::move(node->right));
            return leftRotate(std::move(node));
        }

        return node;
    }

    /*
    Recursive helper for reverse in-order traversal (gets most recent first).
    posts->Vector to store results
    count->Current number of posts found
    */
    void getReverseInOrder(const std::unique_ptr<AVLNode<K, V>>& node, std::vector<V>& posts, int& count, int N) {
        if (!node || (N != -1 && count >= N)) {
            return;
        }

        getReverseInOrder(node->right, posts, count, N);

        if (N == -1 || count < N) {
            posts.push_back(node->value);
            count++;
        }

        getReverseInOrder(node->left, posts, count, N);
    } //Right -> Current -> Left

public:
    AVLTree() : root(nullptr) {}

    void insert(K key, V value) {
        root = insert(std::move(root), key, value);
    }

    std::vector<V> getRecentPosts(int N) {
        std::vector<V> posts;
        int count = 0;
        getReverseInOrder(root, posts, count, N);
        return posts;
    }
};