diff --git a/structure/lazy_segment_tree.cpp b/structure/lazy_segment_tree.cpp new file mode 100644 index 0000000000000000000000000000000000000000..61ecf3991c0e903bdfa18689f992827d4a4f67f9 --- /dev/null +++ b/structure/lazy_segment_tree.cpp @@ -0,0 +1,89 @@ +/** + * Lazy Segment Tree + * + * Complexity (Time): + * build_tree -> O(n log n) + * update_tree -> O(log n) + * query_tree -> O(log n) + * Complexity (Space): O(n) + */ + +int N, tree[2 * MAX], lazy[2 * MAX], v[MAX]; + +// Build tree with elements from v +void build_tree(int node = 1, int a = 0, int b = N) { + if (a > b) + return; + + if (a == b) { + tree[node] = v[a]; + return; + } + + build_tree(node * 2, a, (a + b) / 2); + build_tree(node * 2 + 1, 1 + (a + b) / 2, b); + + tree[node] = tree[node * 2] + tree[node * 2 + 1]; +} + +int modification(int a, int b, int val) { + return ((b - a) + 1) * val; +} + +// Update segment [i,j] by adding value val +void update_tree(int i, int j, int val, int node = 1, int a = 0, int b = N) { + if (lazy[node] != 0) { + tree[node] += modification(a, b, lazy[node]); + + if (a != b) { + lazy[node * 2] += lazy[node]; + lazy[node * 2 + 1] += lazy[node]; + } + + lazy[node] = 0; + } + + if (a > b || a > j || b < i) + return; + + if (a >= i && b <= j) { + tree[node] += modification(a, b, val); + + if (a != b) { + lazy[node * 2] += val; + lazy[node * 2 + 1] += val; + } + + return; + } + + update_tree(i, j, node * 2, a, (a + b) / 2); + update_tree(i, j, node * 2 + 1, 1 + (a + b) / 2, b); + + tree[node] = tree[node * 2] + tree[node * 2 + 1]; +} + +// Return sum from i to j +int query_tree(int i, int j, int node = 1, int a = 0, int b = N) { + if (a > b || a > j || b < i) + return 0; + + if (lazy[node] != 0) { + tree[node] += modification(a, b, lazy[node]); + + if (a != b) { + lazy[node * 2] += lazy[node]; + lazy[node * 2 + 1] += lazy[node]; + } + + lazy[node] = 0; + } + + if (a >= i && b <= j) + return tree[node]; + + int q1 = query_tree(i, j, node * 2, a, (a + b) / 2); + int q2 = query_tree(i, j, node * 2 + 1, 1 + (a + b) / 2, b); + + return q1 + q2; +}