DSA 3

// Gorozhankin Egor BS19-01


public class AVL_Tree<T> implements Comparable<T> {
    private Node<T> top;

    // return height of the node (for null node return 0)
    public static int nodeHeight(Node<?> node) { // O(1)
        if (node == null)
            return 0;
        return node.height;
    }

    // return information about balance of tree or subtree
    // {0, 1, -1} - balanced, otherwise unbalanced
    public static int balanceFactor(Node<?> node) { // O(1)
        if (node == null)
            return 0;
        return nodeHeight(node.left) - nodeHeight(node.right);
    }

    // at field height assigns value of maximum of path to leaf
    public static void calculateNodeHeight(Node<?> node) { // O(1)
        if (nodeHeight(node.left) > nodeHeight(node.right))
            node.height = nodeHeight(node.left) + 1;
        else
            node.height = nodeHeight(node.right) + 1;
    }

    // add data to tree
    public void add(T data) { // O(log n)
        top = add(top, data);
    }

    // add data to tree
    @SuppressWarnings("unchecked")
    private Node<T> add(Node<T> node, T key) { // O(log n)
        if (node == null) {
            return new Node<T>(key);
        }

        if (((Comparable<T>) key).compareTo(node.key) < 0)
            node.left = add(node.left, key);
        else
            node.right = add(node.right, key);

        node.height = 1 + (nodeHeight(node.left) > nodeHeight(node.right) ? nodeHeight(node.left) : nodeHeight(node.right));

        return nodeBalancing(key, node);
    }

    // remove data from tree
    public void remove(T key) { // O(log n)
        top = remove(top, key);
    }

    // remove data from tree
    @SuppressWarnings("unchecked")
    private Node<T> remove(Node<T> node, T key) { // O(log n)
        if (node == null)
            return node;

        if (((Comparable<T>) key).compareTo(node.key) < 0)
            node.left = remove(node.left, key);
        else if (((Comparable<T>) key).compareTo(node.key) > 0)
            node.right = remove(node.right, key);
        else {
            if (node.left == null && node.right == null)
                return null;

            if (node.left == null) {
                Node<T> temp = node.right;
                node = null;
                return temp;
            }

            if (node.right == null) {
                Node<T> temp = node.left;
                node = null;
                return temp;
            }
        }

        Node<T> temp = node.left;
        while (temp.right != null)
            temp = temp.right;
        node.key = temp.key;
        node.left = remove(node.left, temp.key);

        calculateNodeHeight(node);
        return nodeBalancing(node);
    }

    // check balance of tree or tree and make rotations if it need
    private Node<T> nodeBalancing(Node<T> node) { // O(1)
        int balance = balanceFactor(node);

        if (balance > 1) {
            if (balanceFactor(node.left) < 0) {
                node.left = leftRotation(node.left);
            }
            return rightRotation(node);
        }

        if (balance < -1) {
            if (balanceFactor(node.right) > 0) {
                node.right = rightRotation(node.right);
            }
            return leftRotation(node);
        }

        return node;
    }

    // Transformation tree to make it balanced
    private Node<T> rightRotation(Node<T> node) { // O(1)
        Node<T> temp = node.left;
        node.left = temp.right;
        temp.right = node;
        calculateNodeHeight(node);
        calculateNodeHeight(temp);
        return temp;
    }

    // Transformation tree to make it balanced
    private Node<T> leftRotation(Node<T> node) { // O(1)
        Node<T> temp = node.right;
        node.right = temp.left;
        temp.left = node;
        calculateNodeHeight(node);
        calculateNodeHeight(temp);
        return temp;
    }

    // allows us to compare to nodes
    @Override
    public int compareTo(T o) { // O(1)
        return this.compareTo(o);
    }
}

class Node<T> implements Comparable<T> {
    public T key;
    public Node<T> left, right;
    public int height;

    Node(T key) {
        this.key = key;
        left = null;
        right = null;
        height = 1;
    }

    Node(Node<T> node) {
        this.key = node.key;
        this.left = node.left;
        this.right = node.right;
    }

    @Override
    public int compareTo(T o) {
        this.compareTo(o);
        return 0;
    }
}