AVL Tree

index.ts

type AVLNode<T> = {
    left: AVLNode<T>,
    right: AVLNode<T>,
    val: T,
    height: number
} | null;

type CompareResult = -1 | 0 | 1;

function compare<T>(x: T, y: T): CompareResult {
    if (x === y) return 0;
    if (x < y) return -1;
    return 1;``
}

// max
function max(a: number, b: number): number {
    return a > b ? a : b;
}

// abs
function abs(a: number): number {
    return a < 0 ? -1 * a : a;
}

// height of AVL node
function height<T>(n: AVLNode<T>): number {
    if (n === null) return -1;
    // TODO: check if this is needed or not 
    // if (n.left !== null && n.right !== null) return 1 + max(n.left.height, n.right.height);
    // if (n.left === null && n.right !== null) return 1 + n.right.height;
    // if (n.left !== null && n.right === null) return 1 + n.left.height;

    if (n.left !== null && n.right !== null) return n.height;
    if (n.left === null && n.right !== null) return n.height;
    if (n.left !== null && n.right === null) return n.height;

    return 0;
}

// create AVL node
function createAVLNode<T>(x: T) {
    return { left: null, right: null, val: x, height: 0 };
}

// left-rotate
function leftRotate<T>(parent: AVLNode<T>, node: AVLNode<T>): AVLNode<T> {
    if (parent !== null && node !== null) {
        parent.right = node.left;
        node.left = parent;

        parent.height = height(parent);
        node.height = height(node);
    }

    return node;
}

// right-rotate
function rightRotate<T>(parent: AVLNode<T>, node: AVLNode<T>): AVLNode<T> {
    if (parent !== null && node !== null) {
        parent.left = node.right;
        node.right = parent;

        parent.height = height(parent);
        node.height = height(node);
    }

    return node;
}

// balance
function balance<T>(parent: AVLNode<T>, node: AVLNode<T>): AVLNode<T> {
    if (parent !== null && node !== null) {
        const diff = height(parent.left) - height(parent.right);
        if (abs(diff) > 1) {
            if (diff > 0) {
                // left heavy
                if ((height(node.left) - height(node.right)) > 0) {
                    // one rotation
                    return rightRotate(parent, node);
                } else {
                    // two rotations
                    return rightRotate(parent, leftRotate(node, node.right));
                }
            } else {
                // right heavy
                if ((height(node.left) - height(node.right)) < 0) {
                    // one rotation
                    return leftRotate(parent, node);
                } else {
                    // two rotations
                    return leftRotate(parent, rightRotate(node, node.left));
                }
            }
        }
    }
    return parent;
}

// insert
export function insert<T>(root: AVLNode<T>, x: T): AVLNode<T> {
    if (root === null) return createAVLNode(x);

    const cmp = compare(x, root.val);
    if (cmp < 0) {
        root.left = insert(root.left, x);
        root.height = height(root);

        return balance(root, root.left);
    } else if (cmp > 0) {
        root.right = insert(root.right, x);
        root.height = height(root);

        return balance(root, root.right);
    } else {
        return root;
    }
}

// remove
export function remove<T>(root: AVLNode<T>, x: T): AVLNode<T> {
    if (root === null) return null;

    const cmp = compare(x, root.val);
    if (cmp === 0) {
        if (root.left === null && root.right === null) return null;
        if (root.left !== null && root.right === null) return root.left;
        if (root.left === null && root.right !== null) return root.right;

        const maxLeft = findMax(root.left);
        if (maxLeft !== null) {
            root.val = maxLeft.val;
            root.left = remove(root.left, maxLeft.val);
            return balance(root, root.left);
        }


    } else if (cmp < 0) {
        root.left = remove(root.left, x);

        return balance(root, root.left);
    } else {
        root.right = remove(root.right, x);

        return balance(root, root.right);
    }

    return null;
}

// max
export function findMax<T>(root: AVLNode<T>): AVLNode<T> {
    if (root === null) return null;
    if (root.right === null) return root;
    return findMax(root.right)
}

// min

// predecessor

// successor

// inorder
export function inorder<T>(root: AVLNode<T>): T[] {
    const xs: T[] = [];

    function aux(root: AVLNode<T>) {
        if (root !== null) {
            if (root.left !== null) aux(root.left);
            xs.push(root.val);
            if (root.right !== null) aux(root.right);
        }
    }

    aux(root);

    return xs;
}

test.ts

import { insert, remove, findMax, inorder } from '../index';

describe("Test insert node in AVL tree", () => {

    it("should inset node in AVL tree and preserve the AVL property", () => {

        let root = null;
        root = insert(root, 1);
        root = insert(root, 2);
        root = insert(root, 3);
        root = insert(root, 4);
        root = insert(root, 5);
        root = insert(root, 6);
        root = insert(root, 7);
        root = insert(root, 8);
        root = insert(root, 9);
        root = insert(root, 10);
        root = insert(root, 11);
        root = insert(root, 12);
        root = insert(root, 13);
        root = insert(root, 14);
        root = insert(root, 15);

        expect(inorder(root)).toStrictEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]);

    })

    it("should remove node from AVL tree and preserve the AVL property", () => {

        let root = null;
        root = insert(root, 1);
        root = insert(root, 2);
        root = insert(root, 3);
        root = insert(root, 4);
        root = insert(root, 5);

        root = remove(root, 2);
        root = remove(root, 3);

        expect(inorder(root)).toStrictEqual([1, 4, 5]);
    });


})