封裝線段樹實現基本功能

線段樹封裝

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