rank-tree/src/rank.ts

For the structure of sort sets, refer to readme for details

import { RBTreeNode, Pair } from './RB_node';

export class RankTree {
    private _root: RBTreeNode;
    private _hash: Map<string, RBTreeNode>;
    constructor() {
        this._root = null;
        this._hash = new Map<string, RBTreeNode>();
    }

   /**
    * 设置k-v
    * @param key 
    * @param score 
    */
    public tadd(key: string, score: number): void {
        if (this._hash.has(key)) {
            this._onUpdate(key, score);
        } else {
            let node = new RBTreeNode(key, score);
            node.setRed();
            if (!this._root) {
                node.setBlack();
                this._root = node;
                this._hash.set(key, node);
            } else {
                this._onInsert(node, this._root);
            }
        }
    }

    /**
     * 增加某个字段积分
     * @param key 
     * @param add 
     */
    public tincrease(key: string, add: number): void {
        if (this._hash.has(key)) {
            let node = this._hash.get(key);
            let score = node.score + add;
            this._onUpdate(key, score);
        } else {
            let node = new RBTreeNode(key, add);
            if (!this._root) {
                node.setBlack();
                this._root = node;
                this._hash.set(key, node);
            } else {
                this._onInsert(node, this._root);
            }
        }
    }

    /**
     * 扣除指定key值的积分
     * @param key 
     * @param mul 
     */
    public treduce(key: string, mul: number): void {
        if (this._hash.has(key)) {
            let node = this._hash.get(key);
            let score = node.score - mul;
            this._onUpdate(key, score);
        } else {
            let node = new RBTreeNode(key, -mul);
            if (!this._root) {
                node.setBlack();
                this._root = node;
                this._hash.set(key, node);
            } else {
                this._onInsert(node, this._root);
            }
        }
    }

    /**
     * 获取排名
     * @param key 
     * @returns 
     */
    public trank(key: string): number {
        let node = this._hash.get(key);
        if (!node) {
            return -1;
        }
        return this._onGetIndex(this._root, node, 0);
    }

    /**
     * 指定积分的个数
     * @param min 
     * @param max 
     * @returns 
     */
    public tcount(minS: number, maxS: number): number {
        // 定位到起始节点后中序遍历
        return this.trevrangebyscore(minS, maxS).length;
    }

    /**
     * 红黑树节点个数
     * @param min 
     * @param max 
     * @returns 
     */
    public tcard(): number {
        return this._hash.size;
    }

    /**
     * 返回指定序号的元素 升序
     * @param minI 1起始
     * @param maxI 
     * @returns 
     */
    public trange(minI: number, maxI: number = 0): Array<Pair> {
        let _min = Math.max(minI, 1);
        let _max;
        if (maxI == 0) {
            _max = this._hash.size;
        } else {
            _max = Math.min(maxI, this._hash.size);
        }
        let count = _max - _min + 1;
        let list = new Array<Pair>();
        let node = this._findIndexStartNode(this._root, _min);
        if (!node) {
            return list;
        }
        this._LDR_selectIndexRange(node, list, count, _min, _min, _max);
        return list;
    }

    /**
     * 区间积分
     * @param minS 
     * @param maxS 
     * @returns 
     */
    public trevrangebyscore(minS: number, maxS: number = -1): Array<Pair> {
        let list = new Array<Pair>();
        let prev = this._findScoreStartNode(this._root, minS);
        if (!prev) {
            return list;
        }
        if (maxS == -1) {
            maxS = Infinity;
        }
        this._LDR_selectScoreRange(prev, list, minS, maxS);
        return list;
    }

    /**
     * 获取所有数据
     * @returns 
     */
    public tgetall(): Array<Pair> {
        let list = new Array<Pair>();
        this._LDR_SCORE(this._root, list, Infinity);
        return list;
    }

    /**
     * 获取元素
     * @param key 
     * @returns 
     */
    public tscore(key: string): number  {
        let node = this._hash.get(key);
        return node && node.score || 0;
    }

    /**
     * 删除指定key
     * @param key 
     * @returns 
     */
    public trem(keys: string | Array<string>): void {
        let _keys;
        if (!Array.isArray(keys)) {
            _keys = [keys];
        } else {
            _keys = keys;
        }
        _keys.forEach((key) => {
            let node = this._hash.get(key);
            if (!node) {
                return;
            }
            this._onRemove(node)
            this._hash.delete(key);
        })
    }

    /**
     * 获取指定序号元素·
     * @param index 
     * @returns 
     */
    public tgetwithindex(index: number): Pair {
        let node = this._findIndexStartNode(this._root, index);
        if (!node) {
            return null;
        }
        return [node.key, node.score];
    }

    /************************************** 以下为自调用函数 *********************************/
    /**
     * 寻找位置点
     * @param e 
     * @param parent 
     */
    private _onInsert(e: RBTreeNode, parent: RBTreeNode): void {
        if (this._compare(e, parent)) {
            parent.index++;
            if (!parent.left) {
                parent.left = e;
                e.parent = parent;
                this._hash.set(e.key, e);
                this._onTuneUp(e);
            } else {
                parent = parent.left;
                this._onInsert(e, parent);
            }
        } else {
            if (!parent.right) {
                parent.right = e;
                e.parent = parent;
                this._hash.set(e.key, e);
                this._onTuneUp(e);
            } else {
                parent = parent.right;
                this._onInsert(e, parent);
            }
        }
    }

    /**
     * 更新积分处理
     * @param key 
     * @param score 
     */
    private _onUpdate(key: string, score: number): void {
        // 直接删除 重新添加新值， 因为
        let e = this._hash.get(key);
        this._onRemove(e);
        let node = new RBTreeNode(key, score);
        this._hash.set(key, node);
        if (!this._root) {
            node.setBlack();
            this._root = node;
            this._hash.set(key, node);
        } else {
            this._onInsert(node, this._root);
        }
    }

    /**
     * 移除
     * @param node 
     * @returns 
     */
    public _onRemove(node: RBTreeNode): boolean {
        let parent = node.parent;
        if (node.isLeaf()) {
            if (node.isRed()) {
                // 递归调整index
                // 红叶子节点 直接删除
                if (node.isLeftChild()) {
                    parent.left = null;
                    node.parent.index--;
                } else {
                    parent.right = null;
                }
            } else {
                // 先修复 再删除
                this._onRepairNode(node);
                if (node.isLeftChild()) {
                    node.parent.left = null;
                    node.parent.index--;
                } else {
                    node.parent.right = null;
                }
            }
            this._onRepairIndex(node.parent)
        } else {
            // 非叶子置换
            if (node.left && !node.right) {
                // 和左节点置换
                node = this._swapNode(node, node.left);
            } else if (!node.left && node.right) {
                // 和右节点置换
                node = this._swapNode(node, node.right);
            } else {
                // 双节点 寻找后继节点
                let _swapNode = this._successor(node);
                node = this._swapNode(node, _swapNode);
            }
            this._onRemove(node);
        }
        return true;
    }

    /**
     * 修复序号
     * @param node 
     * @returns 
     */
    private _onRepairIndex(node: RBTreeNode): void {
        if (!node.parent) {
            return;
        }
        if (node.isLeftChild()) {
            node.parent.index--;
        }
        this._onRepairIndex(node.parent);
    }

    /**
     * 修复节点， 调整平衡
     * @param node 
     */
    private _onRepairNode(node: RBTreeNode): void {
        if (node.isRoot() || node.isRed()) {
            return;
        }
        let parent = node.parent;
        let brother = node.getBrother();
        if (brother.isBlack()) {
            // 兄弟节点为黑色
            if (node.isLeftChild()) {
                // 左节点
                if (brother.right && !brother.left) {
                    // 只有右节点
                    // 左旋转父节点
                    this._rotateLeft(parent);
                    brother.setColor(parent.color);
                    brother.right.setBlack();
                    parent.setBlack();
                } else if (!brother.right && brother.left) {
                    brother.left.setBlack();
                    brother.setRed();
                    this._rotateRight(brother);
                    this._onRepairNode(node);
                } else if (brother.left && brother.right) {
                    brother.setColor(parent.color);
                    brother.right.setBlack();
                    parent.setBlack();
                    this._rotateLeft(parent);
                } else if (!brother.left && !brother.right) {
                    brother.setRed();
                    this._onRepairNode(parent);
                }
            } else { 
                // 右节点
                if (brother.left && !brother.right) {
                    brother.setColor(parent.color);
                    brother.left.setBlack();
                    parent.setBlack();
                    this._rotateRight(parent);
                } else if (!brother.left && brother.right) {
                    brother.right.setBlack();
                    brother.setRed();
                    this._rotateLeft(brother);
                    this._onRepairNode(node);
                } else if (brother.left && brother.right) {
                    brother.setColor(parent.color);
                    brother.left.setBlack();
                    parent.setBlack();
                    this._rotateRight(parent);
                } else if (!brother.left && !brother.right) {
                    brother.setRed();
                    this._onRepairNode(parent);
                }
            }
        } else {
            // 兄弟节点为红色
            if (node.isLeftChild()) {
                brother.setBlack();
                brother.left.setRed();
                this._rotateLeft(parent);
            } else {
                brother.setBlack();
                brother.right.setRed();
                this._rotateRight(parent);
            }
        }
    }
    /**
     * 调整树 使之保持平衡
     * @param node 节点 
     */
    private _onTuneUp(node: RBTreeNode): void {
        let parent = node.parent;
        if (!parent) {
            // 根节点
            node.setBlack();
            return;
        }
        if (parent.isRed()) {
            let grandparent = parent.parent;
            let uncle = parent.getBrother();
            if (uncle && uncle.isRed()) {
                uncle.setBlack();
                parent.setBlack();
                grandparent.setRed();
                this._onTuneUp(grandparent);
            } else if (!uncle || uncle.isBlack()) {
                let tmp = grandparent.parent;
                let isLeft = grandparent.isLeftChild();
                if (parent.isLeftChild() && node.isLeftChild()) {
                    // 左左
                    grandparent.parent = parent;
                    grandparent.left = parent.right;
                    grandparent.index = (parent.right ? parent.right.index : 0) + 1;
                    if (parent.right) {
                        parent.right.parent = grandparent;
                    }
                    grandparent.setRed();
                    parent.right = grandparent;
                    parent.setBlack();
                    if (!tmp) {
                        this._root = parent;
                        parent.parent = null;
                    } else {
                        if (isLeft) {
                            tmp.left = parent;
                        } else{
                            tmp.right = parent;
                        }
                        parent.parent = tmp;
                    }
                } else if (parent.isLeftChild() && node.isRightChild()) {
                    // 左右
                    grandparent.left = node;
                    node.parent = grandparent;
                    parent.right = node.left;
                    if (node.left) {
                        node.left.parent = parent;
                    }
                    parent.parent = node;
                    node.left = parent;
                    node.index += parent.index;
                    this._onTuneUp(parent);                    
                } else if (parent.isRightChild() && node.isRightChild()) {
                    // 右右
                    grandparent.parent = parent;
                    grandparent.right = parent.left;
                    if (parent.left) {
                        parent.left.parent = grandparent;
                    }
                    grandparent.setRed();
                    parent.index += grandparent.index;
                    parent.left = grandparent;
                    parent.setBlack();
                    if (!tmp) {
                        parent.parent = null;
                        this._root = parent;
                    } else {
                        if (isLeft) {
                            tmp.left = parent;
                        } else{
                            tmp.right = parent;
                        }
                        parent.parent = tmp;
                    }
                } else if (parent.isRightChild() && node.isLeftChild()) {
                    // 右左
                    grandparent.right = node;
                    node.parent = grandparent;
                    parent.left = node.right;
                    if (node.right) {
                        node.right.parent = parent;
                    }
                    parent.parent = node;
                    parent.index -= node.index;
                    node.right = parent;
                    this._onTuneUp(parent);  
                }
            }
        }
    }

    /**
     * 寻找后继节点
     * @param node 
     */
    private _successor(node: RBTreeNode): RBTreeNode {
        if (node == null){
        	return null;
        }
        if (null != node.right) { // 获取 后继节点
        	let p = node.right;
            while (null != p.left){
            	p = p.left;
            }
            return p;
        } else {
        	let p = node.parent;
        	let ch = node;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    /**
     * 左旋转
     * @param node 
     */
    private _rotateLeft(node: RBTreeNode): void {
        let isLeft = node.isLeftChild();
        let parent = node.parent;
        let right = node.right;

        // 更新index
        right.index += node.index;

        node.right = right.left;
        if (right.left) {
            right.left.parent = node;
        }
        // 链接
        right.left = node;
        node.parent = right;

        right.parent = parent; 
        if (parent) {
            if (isLeft) {
                parent.left = right;
            } else {
                parent.right = right;
            }
        } else {
            this._root = right;
        }
    }

    /**
     * 左旋转
     * @param node 
     */
    private _rotateRight(node: RBTreeNode): void {
        let isLeft = node.isLeftChild();
        let parent = node.parent;
        let left = node.left;

        // 更新index
        // left.index += node.index;
        node.index = (left.right ? left.right.index : 0) + 1;
        
        node.left = left.right;
        if (left.right) {
            left.right.parent = node;
        }

        // 链接
        left.right = node;
        node.parent = left;

        left.parent = parent;
        if (parent) {
            if (isLeft) {
                parent.left = left;
            } else {
                parent.right = left;
            }
        } else {
            this._root = left;
        }
    }

    /**
     * 交换节点位置  简单做法为值交换 这里采用节点交换 (弃用)
     * @param nodeFirst 
     * @param nodeSecond 
     */
    private __swapNode(nodeFirst: RBTreeNode, nodeSecond: RBTreeNode) {
        let firstParent = nodeFirst.parent;
        let firstLeft = nodeFirst.left;
        let firstRight = nodeFirst.right;
        let firstIsLeft = nodeFirst.isLeftChild();

        let secondParent = nodeSecond.parent;
        let secondLeft = nodeSecond.left;
        let secondRight = nodeSecond.right;
        let secondIsLeft = nodeSecond.isLeftChild();

        // TODO 交换的节点为父子关系 需要处理
        // 重定向first node
        nodeFirst.parent = secondParent;
        if (secondParent) {
            if (secondIsLeft) {
                secondParent.left = nodeFirst;
            } else {
                secondParent.right = nodeFirst;
            }
        }

        nodeFirst.left = secondLeft;
        if (secondLeft) {
            secondLeft.parent = nodeFirst;
        }
        nodeFirst.right = secondRight;
        if (secondRight) {
            secondRight.parent = nodeFirst;
        }

        // 重定向second node
        nodeSecond.parent = firstParent;
        if (firstParent) {
            if (firstIsLeft) {
                firstParent.left = nodeFirst;
            } else {
                secondParent.right = nodeFirst;
            }
        }
        nodeSecond.left = firstLeft;
        if (firstLeft) {
            firstLeft.parent = nodeSecond;
        }
        nodeSecond.right = firstRight;
        if (firstRight) {
            firstRight.parent = nodeSecond;
        }
    }

    /**
     * 交换节点位置  简单值交换
     * @param nodeFirst 
     * @param nodeSecond 
     */
     private _swapNode(nodeFirst: RBTreeNode, nodeSecond: RBTreeNode) {
        let keyFirst = nodeFirst.key;
        let scoreFirst = nodeFirst.score;
        let keySecond = nodeSecond.key;
        let scoreSecond = nodeSecond.score;

        nodeFirst.key = keySecond;
        nodeFirst.score = scoreSecond;
        nodeSecond.key = keyFirst;
        nodeSecond.score = scoreFirst;
        this._hash.set(nodeFirst.key, nodeFirst);
        this._hash.set(nodeSecond.key, nodeSecond);
        return nodeSecond;
    }
    
    /**
     * 比较器
     * @param e 当前
     * @param target 目标 
     * @returns 
     */
    private _compare(e: RBTreeNode, target: RBTreeNode): boolean {
        return e.score < target.score;
    }

    /**
     * 获取序号
     * @param prev 
     * @param node 
     * @param index 
     * @returns 
     */
    private _onGetIndex(prev: RBTreeNode, node: RBTreeNode, index: number): number {
        if (!prev) {
            return -1;
        }
        if (node == prev) {
            return prev.index + index;
        }
        if (this._compare(node, prev)) {
            return this._onGetIndex(prev.left, node, index);         
        } else {
            return this._onGetIndex(prev.right, node, index+prev.index);
        }
    }

    /**
     * 寻找指定分数的前驱节点 
     * @param prev 
     * @param min 
     * @returns 
     */
    private _findScoreStartNode(prev: RBTreeNode, min: number): RBTreeNode {
        if (!prev) {
            return null;
        }
        if (min <= prev.score) {
            if (!prev.left) {
                return prev;
            } else {
                let res = this._findScoreStartNode(prev.left, min);
                if (res) {
                    return res;
                }
                return prev;
            }
        } else {
            if (!prev.right) {
                return null;
            } else {
                return this._findScoreStartNode(prev.right, min);
            }
        }
    }

    private _LDR_selectScoreRange(prev: RBTreeNode, list: Array<Pair>, min: number, max: number): void {
        if (!prev) {
            return;
        }
        if (prev.score < min) {
            this._LDR_selectScoreRange(prev.parent, list, min, max);
            return;
        }
        // 先解决此节点下的问元素
        if (prev.score > max) {
            return;
        }
        list.push([prev.key, prev.score]);
        // 解决右子树
        this._LDR_SCORE(prev.right, list, max);
        this._LDR_selectScoreRange(prev.parent, list, min, max);
    }

    private _LDR_SCORE(prev: RBTreeNode, list: Array<Pair>, max: number): void {
        if (!prev) {
            return;
        }
        this._LDR_SCORE(prev.left, list, max);
        if (prev.score > max) {
            return;
        }
        list.push([prev.key, prev.score]);
        this._LDR_SCORE(prev.right, list, max);
    }

    private _findIndexStartNode(prev: RBTreeNode, minI: number): RBTreeNode {
        if (!prev) {
            return null;
        }
        if (prev.index == minI) {
            return prev;
        } else if (prev.index < minI) {
            minI -= prev.index;
            return this._findIndexStartNode(prev.right, minI);
        }
        return this._findIndexStartNode(prev.left, minI);
    }

    private _LDR_selectIndexRange(prev: RBTreeNode, list: Array<Pair>, count: number, prevI: number, minI: number, maxI: number): void {
        if (!prev) {
            return;
        }

        // 已经满足筛选
        if (list.length >= count) {
            return;
        }
        let nextPrevI;
        // 计算父节点 prevI
        if (prev.parent) {
            if (prev.isLeftChild()) {
                nextPrevI = prevI + (prev.parent.index - prev.index);
            } else {
                nextPrevI = prevI - prev.index;
            }
        }
        
        if (prevI < minI) {
            if (prev.parent) {
                this._LDR_selectIndexRange(prev.parent, list, count, nextPrevI, minI, maxI);
            }
            return;
        }
        list.push([prev.key, prev.score]);
        // 解决右子树
        this._LDR_INDEX(prev.right, list, count);
        if (prev.parent) {
            this._LDR_selectIndexRange(prev.parent, list, count, nextPrevI, minI, maxI);
        }
    }

    private _LDR_INDEX(prev: RBTreeNode, list: Array<Pair>, count: number): number {
        if (!prev) {
            return; 
        }
        this._LDR_INDEX(prev.left, list, count);
        // 满足筛选
        if (list.length >= count) {
            return;
        }
        list.push([prev.key, prev.score]);
        this._LDR_INDEX(prev.right, list, count);
    }
}