competitive-programming/segment_tree.cpp at master · InsanePrototyper/competitive-programming · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class SegmentTree {
public:
    SegmentTree *lchild, *rchild;
    int leftindex, rightindex;
    ll val = 0;
    ll lazy = 0;
    ll combine() {
        if(leftindex == rightindex) return this->val;
        return (lchild->val + rchild->val); // change
    }

    void updateLazy() {
        val += (rightindex - leftindex + 1) * lazy;
        if(leftindex != rightindex) {
            lchild->lazy += lazy;
            rchild->lazy += lazy;
        }
        lazy = 0;
    }

    SegmentTree(){}
    SegmentTree(int leftindex, int rightindex, std::vector<ll> &v) {
        this->leftindex = leftindex;
        this->rightindex = rightindex;
        if(leftindex == rightindex) {
            val = v[leftindex]; // change
            return;
        }
        int mid = (leftindex + rightindex) >> 1;
        lchild = new SegmentTree(leftindex, mid, v);
        rchild = new SegmentTree(mid + 1, rightindex, v);
        this->val = combine();
    }

    void pUpdate(int pos, ll val) {
        updateLazy();
        if(pos < leftindex or pos > rightindex) {
            return;
        }
        if(leftindex == rightindex) {
            this->val = val; // change
            return;
        }
        if(pos <= lchild->rightindex)
            lchild->pUpdate(pos, val);
        else
            rchild->pUpdate(pos, val);
        this->val = combine();
    }

    ll pQuery(int pos) {
        updateLazy();
        if(pos < leftindex || pos > rightindex) {
            return 0;
        }
        if(leftindex == rightindex) {
            return val;
        }
        if(pos <= lchild->rightindex) {
            return lchild->pQuery(pos);
        }
        else {
            return rchild->pQuery(pos);
        }
    }

    ll rQuery(int l, int r) {
        updateLazy();
        if(r < leftindex || l > rightindex) {
            return 0; // change
        }
        if(l <= leftindex && r >= rightindex) {
            return val;  // change
        }
        return lchild->rQuery(l, r) + rchild->rQuery(l, r);
    }

    void rUpdate(int l, int r, ll val) { //  add val to each element from l to r
        updateLazy();
        if(r < leftindex || l > rightindex) {
            return;
        }
        if(l <= leftindex && r >= rightindex) {
            lazy += val;
            updateLazy();
            return;
        }
        lchild->rUpdate(l, r, val);
        rchild->rUpdate(l, r, val);
        this->val = combine();
    }

};