Untitled

//
// Created by erikr on 5/10/2021.
//

#pragma once

#include <memory>

template <typename Key, typename Value=Key>
class AVLTree {
public:
    class Node {
    private:
        Key key_;
        Value value_;
        int balance_factor_, height_;
        Node *left_, *right_;
        Node(const Key& key) : key_(key), height_(0) {}
        Node(const Key& key, const Value& value) : key_(key), value_(value), height_(0) {}

    public:
        // All these getter methods are O(1)
        Node *left() {
            return left_;
        }
        Node *right() {
            return right_;
        }
        const Key& key() const {
            return key_;
        }
        const Value& value() const {
            return value_;
        }
        const int balance_factor() const {
            return balance_factor_;
        }

        friend class AVLTree<Key, Value>;
    };
// class AVLTree
private:
    Node *root_;
    int size_;

    void insert( const Key& new_key, const Value& new_value)
    {
        Node *cur_node = root_;

        if( root_ == nullptr ) {
            root_ = new Node(new_key, new_value);
            size_++;
            balance(cur_node);
            return;
        }

        while (true)
        {
            if( new_value < cur_node->value_ )
            {
                if (cur_node->left_ == nullptr) {
                    cur_node->left_ = new Node(new_key,new_value);
                    size_++;
                    balance(cur_node);
                    return;
                } else {
                    cur_node = cur_node->left_;
                }
            }

            else if( new_value > cur_node->value_ )
            {
                if (cur_node->right_ == nullptr){
                    cur_node->right_ = new Node(new_value,new_key);
                    size_++;
                    balance(cur_node);
                    return;
                } else {
                    cur_node = cur_node->right_;
                }
            }
        }
    }

public:
    AVLTree() : size_(0), root_(nullptr){}

    // TODO: Add code to update the balance factor and do rebalancing...

    static const int ALLOWED_IMBALANCE = 1;

    void insert(const Value &x)
    {
        insert(x,root_);
    }

    void balance(Node *&rootOfSubtree)
    {
        if ( rootOfSubtree == nullptr )
            return;

        if ( height( rootOfSubtree->left_ ) - height( rootOfSubtree->right_ ) > ALLOWED_IMBALANCE )
            if (height( rootOfSubtree->left_->left_ ) >= height(rootOfSubtree->left_->right_))
                rotateWithLeftChild(rootOfSubtree);

    }
    // This overloaded operator method is O(log(size()))
    Value& operator[](const Key& key) {
        // Try to find the node with the value we want:
        Node *cur;
        for (cur = root_;
             cur != nullptr;
             cur = key < (*cur).key_ ? &(*cur).left_ : &(*cur).right_) {
            if (key == (*cur).key_) {
                return (*cur).value_;
            }
        }
        // If we didn't find it, insert a new node with that key:
        // (This is the same behaviour as an std::map.)
        cur = new Node(key);
        ++size_;
        return (*cur).value_;
    }

    int size() {
        return size_;
    }

    Node *root() {
        return root_;
    }

    int height( Node *rootOfSubtree ) const
    {
        return rootOfSubtree == nullptr ? -1 : rootOfSubtree->height_;
    }

    int max(int lhs, int rhs) const
    {
        return lhs > rhs ? lhs : rhs;
    }

    //
    // ROTATIONS
    //

    void rotateWithLeftChild(Node *&k2)
    {
        Node *k1 = k2->left_;
        k2->left_ = k1->right_;
        k1->right_ = k2;
        k2->height_ = max(height( k2->left_ ), height( k2->right_ ) ) + 1;
        k1->height_ = max(height( k1->left_), k2->height_ ) + 1;
        k2 = k1;
    }

    void rotateWithRightChild(Node *&k1)
    {
        Node *k2 = k1->right_;
        k1->right = k2->left_;
        k2->left_ = k1;
        k1->height_ = max(height( k1->left_ ), height( k1->right_)) + 1;
        k2->height_ = max(height( k2->right_), k1->height_) + 1;
        k1 = k2;
    }

    void doubleWithLeftChild(Node *&k3)
    {
        rotateWithRightChild(k3->left_);
        rotateWithLeftChild(k3);
    }

    void doubleWithRightChild(Node *&k1)
    {
        rotateWithLeftChild(k1->right_);
        rotateWithRightChild(k1);
    }

};