diff --git a/structure/segment_tree.cpp b/structure/segment_tree.cpp new file mode 100644 index 0000000000000000000000000000000000000000..a78608dc241c48714bf5da7f9f11c800dd86bc8c --- /dev/null +++ b/structure/segment_tree.cpp @@ -0,0 +1,56 @@ +/** + * Segment Tree + * Complexity (Time): + * Build -> O(n log n) + * Update -> O(log n) + * Query -> O(log n) + * Complexity (Space): O(n) + */ + +int N, tree[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]; +} + +// Update position idx with value val +void update_tree(int idx, int val, int node = 1, int a = 0, int b = N) { + if (a > b || a > idx || b < idx) + return; + + if (a == b) { + tree[node] = val; + return; + } + + update_tree(idx, val, node * 2, a, (a + b) / 2); + update_tree(idx, val, 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 (a >= i && b <= j) + return tree[node]; + + int res1 = query_tree(i, j, node * 2, a, (a + b) / 2); + int res2 = query_tree(i, j, node * 2 + 1, 1 + (a + b) / 2, b); + + return res1 + res2; +}