封装线段树实现基本功能

线段树封装

1. 结点类TreeNode

/* 树结点*/
function TreeNode(x, y, val) {
    this.x = x;
    this.y = y;
    this.val = val || null;
}

2. 线段树类

通过options可以传入pushUp函数设置父子结点关系,root记录线段树根节点

function SegTree(options) {
    this.left = null;
    this.right = null;
    this.nums = options.nums;//数组
    this.pushUp = options.pushUp;//自定义父子关联的函数
    this.root = this.build(0, this.nums.length - 1);//建立线段树
}

3. 建树

SegTree.prototype.build = function (x, y) {
    if (x === y) {/* 创建叶子结点并返回结点 */
        return new TreeNode(x, y, this.nums[x]);
    }
    /* 创建结点 */
    let node = new TreeNode(x, y);
    /* 创建子结点 */
    let mid = Math.floor(x + y >> 1);
    node.left = this.build(x, mid);
    node.right = this.build(mid + 1, y);
    /* 关联父子结点 */
    this.pushUp(node, node.left, node.right);
    return node;
}

4. 查询

SegTree.prototype.query = function (tgtX, tgtY, node) {
    //默认为root
    node = node || this.root;
    let x = node.x, y = node.y;
    /* 在[tgtX,tgtY]区域内 */
    if (x >= tgtX && y <= tgtY) {
        return node.val;
    }
    /* 将查询到的值进行叠加 */
    let count = 0;
    let m = Math.floor(x + y >> 1);
    if (tgtY <= m)
        count += this.query(tgtX, tgtY, node.left);
    if (tgtX > m)
        count += this.query(tgtX, tgtY, node.right);
    return count;
}

5. 更新

SegTree.prototype.update = function (tgtX, tgtY, val, node) {
    //默认为root
    node = node || this.root;
    if (node.x === node.y) {
        node.val += val;
        return;
    }
    let mid = Math.floor(tgtX + tgtY >> 1);
    if (tgtX <= mid)
        this.update(tgtX, tgtY, val, node.left)
    if (tgtY > mid)
        this.update(tgtX, tgtY, val, node.right);
    /* 更新父节点 */
    this.pushUp(node, node.left, node.right);
}

6. 删除

SegTree.prototype.delete = function (tgtX, tgtY, node) {
    //默认为root
    node = node || this.root;
    if (tgtX >= node.x && tgtY <= node.y) {
        node.val = 0;
        return;
    }
    let mid = Math.floor(tgtX + tgtY >> 1);
    let flagx, flagy;//标识左右结点是否删除
    if (tgtX <= mid) {
        this.delete(tgtX, tgtY, node.left);
        flagx = true;
    }
    if (tgtY > mid) {
        this.delete(tgtX, tgtY, node.right);
        flagy = true;
    }
    /* 更新父节点 */
    this.pushUp(node, node.left, node.right);
    if (flagx) node.left = null;
    if (flagy) node.right = null;
}

测试代码

/* 测试 */
let tree = new SegTree({
    nums: [-2, 1, -3, 4, -1, 2, 1, -5, 4],
    pushUp: (node, xNode, yNode) => {
        node.val = xNode.val + yNode.val;
    }
})
console.log(tree.query(0, 8))//1
tree.update(0, 8, 1)
console.log(tree.query(0, 8))//10