CP-codes/Data Structures/segtree.cpp at master · hemant2132/CP-codes · 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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
/*
    -> Segment Tree

    -> ref:
        1. https://www.hackerearth.com/practice/data-structures/advanced-data-structures/segment-trees/tutorial/
        2. https://cp-algorithms.com/data_structures/segment_tree.html
        3. https://www.hackerearth.com/practice/notes/segment-tree-and-lazy-propagation/

    -> Implementation for:
        1. point update and range query
        2. range update and point query
        3. range update and range query (lazy propagation)
*/

int segtree[4 * N];
int lazy[4 * N];

void build(int node, int beg, int en) {
    if (beg == en) {
        // segtree[node] = base[beg];
        return;
    }

    int mid = (beg + en) / 2;
    build(2 * node, beg, mid);
    build(2 * node + 1, mid + 1, en);

    segtree[node] = segtree[2 * node] + segtree[2 * node + 1];
}

// for point update and range query:
void pointUpdate(int node, int beg, int en, int idx, int val) {
    if (beg == en) {
        // base[idx] = val;
        segtree[node] = val;
        return;
    }

    int mid = (beg + en) / 2;
    if (beg <= idx && idx <= mid) {
        pointUpdate(2 * node, beg, mid, idx, val);
    } else {
        pointUpdate(2 * node + 1, mid + 1, en, idx, val);
    }

    segtree[node] = segtree[2 * node] + segtree[2 * node + 1];
}

int rangeQuery(int node, int beg, int en, int l, int r) {
    if (r < beg || en < l) return 0;
    if (l <= beg && en <= r) return segtree[node];

    int mid = (beg + en) / 2;
    int q1 = rangeQuery(2 * node, beg, mid, l, r);
    int q2 = rangeQuery(2 * node + 1, mid + 1, en, l, r);

    return q1 + q2;
}

// for range update and point query:
void rangeUpdate(int node, int beg, int en, int l, int r, int val) {
    if (r < beg || en < l) return;

    if (l <= beg && en <= r) {
        segtree[node] += val;
        return;
    }

    int mid = (beg + en) / 2;
    rangeUpdate(2 * node, beg, mid, l, r, val);
    rangeUpdate(2 * node + 1, mid + 1, en, l, r, val);
}

int pointQuery(int node, int beg, int en, int idx, int val = 0) {
    if (beg == en) return val + segtree[node];

    int mid = (beg + en) / 2;
    if (beg <= idx && idx <= mid) {
        return pointQuery(2 * node, beg, mid, idx, val + segtree[node]);
    }

    return pointQuery(2 * node + 1, mid + 1, en, idx, val + segtree[node]);
}

// for lazy update and lazy query:
void lazyUpdate(int node, int beg, int en, int l, int r, int val) {
    if (lazy[node] != 0) {
        segtree[node] += (en - beg + 1) * lazy[node];
        if (beg != en) {
            lazy[2 * node] += lazy[node];
            lazy[2 * node + 1] += lazy[node];
        }
        lazy[node] = 0;
    }

    if (beg > en || r < beg || en < l) return;

    if (l <= beg && en <= r) {
        segtree[node] += (en - beg + 1) * val;
        if (beg != en) {
            lazy[2 * node] += val;
            lazy[2 * node + 1] += val;
        }

        return;
    }

    int mid = (beg + en) / 2;
    lazyUpdate(2 * node, beg, mid, l, r, val);
    lazyUpdate(2 * node + 1, mid + 1, en, l, r, val);

    segtree[node] = segtree[2 * node] + segtree[2 * node + 1];
}

int lazyQuery(int node, int beg, int en, int l, int r) {
    if (beg > en || r < beg || en < l) return 0;

    if (lazy[node] != 0) {
        segtree[node] += (en - beg + 1) * lazy[node];
        if (beg != en) {
            lazy[2 * node] += lazy[node];
            lazy[2 * node + 1] += lazy[node];
        }
        lazy[node] = 0;
    }

    if (l <= beg && en <= r) return segtree[node];

    int mid = (beg + en) / 2;
    int q1 = lazyQuery(2 * node, beg, mid, l, r);
    int q2 = lazyQuery(2 * node + 1, mid + 1, en, l, r);

    return q1 + q2;
}