AVL木辞書（1つのキーに複数の値を含めることができます）

AVL木辞書とは

AVL木辞書はAVL木の拡張であり、1つのインデックスに複数の値を格納することができます.

データ構造

1つノードのデータ構造,1つノードに複数の値を格納することができます.

class N<T> {
    key: number;
    values = new Array<T>();
    height = 1;
    left: N<T> | undefined;
    right: N<T> | undefined;
}

全部のコード

import * as assert from "typed-assert";

class N<T> {
    //NOTE Key will be set when remove node
    key: number;
    values = new Array<T>();
    height = 1;
    left: N<T> | undefined;
    right: N<T> | undefined;

    get value(): T {
        assert.isNotUndefined(this.values[0]);
        return this.values[0];
    }

    get count() {
        return this.values.length;
    }

    constructor(key: number, value: T) {
        this.key = key;
        this.values.push(value);
    }

    addRepeat(value: T) {
        this.values.push(value);
    }

    removeRepeat(value: T) {
        const index = this.values.indexOf(value);
        if (index > -1) {
            this.values.splice(index, 1);
        }
    }
}

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

export class AvlTree<T> {

    private _root: N<T> | undefined;

    private getHeight(node: N<T> | undefined): number {
        if (node === undefined) {
            return 0;
        }
        return node.height;
    }

    private rightRotate(y: N<T>): N<T> {

        assert.isNotUndefined(y.left);
        const x = y.left;
        const t2 = x.right;

        x.right = y;
        y.left = t2;

        y.height = max(
            this.getHeight(y.left),
            this.getHeight(y.right)) +
            1;
        x.height = max(
            this.getHeight(x.left),
            this.getHeight(x.right)) +
            1;

        return x;
    }

    private leftRotate(x: N<T>): N<T> {
        assert.isNotUndefined(x.right);
        const y = x.right;
        const t2 = y.left;

        y.left = x;
        x.right = t2;

        x.height = max(
            this.getHeight(x.left),
            this.getHeight(x.right)) +
            1;
        y.height = max(
            this.getHeight(y.left),
            this.getHeight(y.right)) +
            1;

        return y;
    }

    private getBalance(node: N<T> | undefined): number {
        if (node === undefined) {
            return 0;
        }
        return this.getHeight(node.left) - this.getHeight(node.right);
    }


    private insertInternal(node: N<T> | undefined, key: number, value: T): N<T> {
        if (node === undefined)
            return (new N(key, value));

        if (key < node.key) {
            node.left = this.insertInternal(node.left, key, value);
        } else if (key > node.key) {
            node.right = this.insertInternal(node.right, key, value);
        } else {
            //NOTE The same key contains different values
            node.addRepeat(value);
            return node;
        }
        node.height = 1 +
            max(
                this.getHeight(node.left),
                this.getHeight(node.right));
        const balance = this.getBalance(node);

        //LL
        if (balance > 1
            && node.left !== undefined
            && key < node.left.key) {
            return this.rightRotate(node);
        }

        //RR
        if (balance < -1
            && node.right !== undefined
            && key > node.right.key) {
            return this.leftRotate(node);
        }

        //LR
        if (balance > 1
            && node.left !== undefined
            && key > node.left.key) {
            node.left = this.leftRotate(node.left);
            return this.rightRotate(node);
        }

        //RL
        if (balance < -1
            && node.right !== undefined
            && key < node.right.key) {
            node.right = this.rightRotate(node.right);
            return this.leftRotate(node);
        }

        return node;
    }

    private getRangeKeysInternal(node: N<T> | undefined, low: number, high: number, array: Array<number>) {
        if (node === undefined) {
            return;
        }

        if (node.key > low && node.key < high) {
            array.push(node.key);
            this.getRangeKeysInternal(node.left, low, high, array);
            this.getRangeKeysInternal(node.right, low, high, array);
        } else if (node.key < low) {
            this.getRangeKeysInternal(node.right, low, high, array);
        } else if (node.key > high) {
            this.getRangeKeysInternal(node.left, low, high, array);
        }
    }

    private getRangeValuesInternal(node: N<T> | undefined, low: number, high: number, array: Array<T>) {
        if (node === undefined) {
            return;
        }
        if (node.key > low && node.key < high) {
            for (let i = 0; i < node.values.length; ++i) {
                array.push(node.values[i]);
            }
            this.getRangeValuesInternal(node.left, low, high, array);
            this.getRangeValuesInternal(node.right, low, high, array);
        } else if (node.key < low) {
            this.getRangeValuesInternal(node.right, low, high, array);
        } else if (node.key > high) {
            this.getRangeValuesInternal(node.left, low, high, array);
        }
    }


    private getMaxNodeInternal(n: N<T>): N<T> {
        while (n.right !== undefined) {
            n = n.right;
        }
        return n;
    }

    private getMinNodeInternal(n: N<T>): N<T> {
        while (n.left !== undefined) {
            n = n.left;
        }
        return n;
    }

    private removeInternal(node: N<T> | undefined, key: number): N<T> | undefined {
        if (node === undefined)
            return node;
        if (key < node.key)
            node.left = this.removeInternal(node.left, key);
        else if (key > node.key)
            node.right = this.removeInternal(node.right, key);
        else {

            if ((node.left === undefined)
                || (node.right === undefined)) {

                let temp: N<T> | undefined;

                if (temp === node.left)
                    temp = node.right;
                else
                    temp = node.left;

                //NOTE No child case 
                if (temp === undefined) {
                    node = undefined;
                } else
                    node = temp;
            } else {
                // Node with two children: Get the inorder 
                // successor (smallest in the right subtree) 
                const temp = this.getMinNodeInternal(node.right);
                // Copy the inorder successor's data to this node 
                node.key = temp.key;
                // Delete the inorder successor 
                node.right = this.removeInternal(node.right, temp.key);
            }
        }

        if (node === undefined)
            return node;

        node.height = max(
            this.getHeight(node.left),
            this.getHeight(node.right)) +
            1;

        const balance = this.getBalance(node);

        //LL
        if (balance > 1 && this.getBalance(node.left) >= 0)
            return this.rightRotate(node);

        //LR
        if (balance > 1 && this.getBalance(node.left) < 0) {
            assert.isNotUndefined(node.left);
            node.left = this.leftRotate(node.left);
            return this.rightRotate(node);
        }

        //RR
        if (balance < -1 && this.getBalance(node.right) <= 0)
            return this.leftRotate(node);

        //RL
        if (balance < -1 && this.getBalance(node.right) > 0) {
            assert.isNotUndefined(node.right);
            node.right = this.rightRotate(node.right);
            return this.leftRotate(node);
        }

        return node;
    }

    insert(key: number, value: T) {
        this._root = this.insertInternal(this._root, key, value);
    }

    getMin(): number {
        return this.getMinNode().key;
    }


    getMinNode(): N<T> {
        assert.isNotUndefined(this._root);
        return this.getMinNodeInternal(this._root);
    }

    getMaxNode(): N<T> {
        assert.isNotUndefined(this._root);
        return this.getMaxNodeInternal(this._root);
    }

    getMax(): number {
        return this.getMaxNode().key;
    }

    getRangeKeys(low: number, high: number): Array<number> {
        const array = new Array<number>();
        this.getRangeKeysInternal(this._root, low, high, array);
        return array;
    }

    getRangeValues(low: number, high: number): Array<T> {
        const array = new Array<T>();
        this.getRangeValuesInternal(this._root, low, high, array);
        return array;
    }

    remove(key: number) {
        this._root = this.removeInternal(this._root, key);
    }


    removeByValue(target: T) {

        if (this._root === undefined) {
            return;
        }

        this.removeByValueInternal(target, this._root);
    }

    private removeByValueInternal(target: T, node: N<T> | undefined) {
        if (node === undefined) {
            return;
        }

        const values = node.values;
        if (values.indexOf(target) !== -1) {

            if (values.length > 1) {
                node.removeRepeat(target);
            } else {
                this.remove(node.key);
                //NOTE Only remove once
                return;
            }

        } else {
            this.removeByValueInternal(target, node.left);
            this.removeByValueInternal(target, node.right);
        }
    }
}