cosmos/code/data_structures/src/segment_tree/segment_tree_lazy.cpp at 96ce61c61829aee0509476711acebd8797e26570 · OpenGenus/cosmos · 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
#include <iostream>
#include <vector>

class SegmentTreeLazy {
public:
    SegmentTreeLazy(const std::vector<int>& data);

    void updateRange(int l, int r, int val);
    int queryRange(int l, int r);

private:
    void build(int node, int start, int end);
    void updateRangeUtil(int node, int start, int end, int l, int r, int val);
    int queryRangeUtil(int node, int start, int end, int l, int r);

    int size_;
    std::vector<int> tree_;
    std::vector<int> lazy_;
    std::vector<int> data_;
};

SegmentTreeLazy::SegmentTreeLazy(const std::vector<int>& data)
    : size_{static_cast<int>(data.size())},
      tree_(4 * size_),
      lazy_(4 * size_, 0),
      data_{data} {
    build(1, 0, size_ - 1);
}

void SegmentTreeLazy::build(int node, int start, int end) {
    if (start == end) {
        tree_[node] = data_[start];
    } else {
        int mid = (start + end) / 2;
        build(2 * node, start, mid);
        build(2 * node + 1, mid + 1, end);
        tree_[node] = tree_[2 * node] + tree_[2 * node + 1];
    }
}

void SegmentTreeLazy::updateRange(int l, int r, int val) {
    updateRangeUtil(1, 0, size_ - 1, l, r, val);
}

void SegmentTreeLazy::updateRangeUtil(int node, int start, int end, int l, int r, int val) {
    if (lazy_[node] != 0) {
        tree_[node] += (end - start + 1) * lazy_[node];
        if (start != end) {
            lazy_[2 * node] += lazy_[node];
            lazy_[2 * node + 1] += lazy_[node];
        }
        lazy_[node] = 0;
    }

    if (start > end || start > r || end < l) {
        return;
    }

    if (start >= l && end <= r) {
        tree_[node] += (end - start + 1) * val;
        if (start != end) {
            lazy_[2 * node] += val;
            lazy_[2 * node + 1] += val;
        }
        return;
    }

    int mid = (start + end) / 2;
    updateRangeUtil(2 * node, start, mid, l, r, val);
    updateRangeUtil(2 * node + 1, mid + 1, end, l, r, val);
    tree_[node] = tree_[2 * node] + tree_[2 * node + 1];
}

int SegmentTreeLazy::queryRange(int l, int r) {
    return queryRangeUtil(1, 0, size_ - 1, l, r);
}

int SegmentTreeLazy::queryRangeUtil(int node, int start, int end, int l, int r) {
    if (start > end || start > r || end < l) {
        return 0;
    }

    if (lazy_[node] != 0) {
        tree_[node] += (end - start + 1) * lazy_[node];
        if (start != end) {
            lazy_[2 * node] += lazy_[node];
            lazy_[2 * node + 1] += lazy_[node];
        }
        lazy_[node] = 0;
    }

    if (start >= l && end <= r) {
        return tree_[node];
    }

    int mid = (start + end) / 2;
    int p1 = queryRangeUtil(2 * node, start, mid, l, r);
    int p2 = queryRangeUtil(2 * node + 1, mid + 1, end, l, r);
    return p1 + p2;
}

int main() {
    std::vector<int> data{1, 3, 5, 7, 9, 11};
    SegmentTreeLazy st(data);

    std::cout << "Sum [1, 3]: " << st.queryRange(1, 3) << "\n";
    st.updateRange(1, 5, 10);
    std::cout << "Sum [1, 3] after update: " << st.queryRange(1, 3) << "\n";

    return 0;
}