|
| 1 | +/** |
| 2 | + * @file |
| 3 | + * @brief A data structure to quickly do operations on ranges: the [Segment Tree](https://en.wikipedia.org/wiki/Segment_tree) algorithm implementation |
| 4 | + * @details |
| 5 | + * Implementation of the segment tree data structre |
| 6 | + * |
| 7 | + * Can do point updates (updates the value of some position) |
| 8 | + * and range queries, where it gives the value of some associative |
| 9 | + * opperation done on a range |
| 10 | + * |
| 11 | + * Both of these operations take O(log N) time |
| 12 | + * @author [Nishant Chatterjee](https://github.com/nishantc1527) |
| 13 | + */ |
| 14 | + |
| 15 | +#include <iostream> /// For IO operations |
| 16 | +#include <vector> /// For std::vector |
| 17 | +#include <algorithm> /// For std::min and std::max |
| 18 | +#include <cassert> /// For assert |
| 19 | + |
| 20 | +/* |
| 21 | + * @namespace |
| 22 | + * @brief Data structures |
| 23 | + */ |
| 24 | +namespace data_structures { |
| 25 | +/** |
| 26 | + * @brief class representation of the segment tree |
| 27 | + * @tparam T The type of the class that goes in the datastructure |
| 28 | + */ |
| 29 | +template <class T> |
| 30 | +class SegmentTree { |
| 31 | +private: |
| 32 | + const T ID = 0; ///< Comb(ID, x) = x |
| 33 | + std::vector<T> t; ///< Vector to represent the tree |
| 34 | + int size = 0; ///< Number of elements available for querying in the tree |
| 35 | +private: |
| 36 | + /** |
| 37 | + * @brief Any associative function that combines x and y |
| 38 | + * @param x The first operand |
| 39 | + * @param y The second operand |
| 40 | + * @return Some associative operation applied to these two values. In this case, I used addition |
| 41 | + */ |
| 42 | + T comb(T x, T y) { |
| 43 | + return x + y; |
| 44 | + } |
| 45 | + /** |
| 46 | + * @brief Gives the midpoint between two integers |
| 47 | + * @param l The left endpoint |
| 48 | + * @param r The right endpoint |
| 49 | + * @return the middle point between them |
| 50 | + */ |
| 51 | + int mid(int l, int r) { |
| 52 | + return l + (r - l) / 2; |
| 53 | + } |
| 54 | + /** |
| 55 | + * @brief Helper method for update method below |
| 56 | + * @param i The index of the current node |
| 57 | + * @param l The leftmost node of the current node |
| 58 | + * @param r The rightmost node of the current node |
| 59 | + * @param pos The position to update |
| 60 | + * @param val The value to update it to |
| 61 | + */ |
| 62 | + void update(int i, int l, int r, int pos, T val) { |
| 63 | + if(l == r) t[i] = val; |
| 64 | + else { |
| 65 | + int m = mid(l, r); |
| 66 | + if(pos <= m) update(i * 2, l, m, pos, val); |
| 67 | + else update(i * 2 + 1, m + 1, r, pos, val); |
| 68 | + t[i] = comb(t[i * 2], t[i * 2 + 1]); |
| 69 | + } |
| 70 | + } |
| 71 | + /** |
| 72 | + * @brief Helper method for range_comb method below |
| 73 | + * @param i The current node |
| 74 | + * @param l The leftmost node of the current node |
| 75 | + * @param r The rightmost node of the current node |
| 76 | + * @param tl The left endpoint of the range |
| 77 | + * @param tr The right endpoint of the range |
| 78 | + * @return The comb operation applied to all values between tl and tr |
| 79 | + */ |
| 80 | + T range_comb(int i, int l, int r, int tl, int tr) { |
| 81 | + if(l == tl && r == tr) return t[i]; |
| 82 | + if(tl > tr) return 0; |
| 83 | + int m = mid(l, r); |
| 84 | + return comb(range_comb(i * 2, l, m, tl, std::min(tr, m)), range_comb(i * 2 + 1, m + 1, r, std::max(tl, m + 1), tr)); |
| 85 | + } |
| 86 | +public: |
| 87 | + SegmentTree(int n) : t(n * 4, ID), size(n) {} |
| 88 | + /** |
| 89 | + * @brief Updates a value at a certain position |
| 90 | + * @param pos The position to update |
| 91 | + * @param val The value to update it to |
| 92 | + */ |
| 93 | + void update(int pos, T val) { |
| 94 | + update(1, 1, size, pos, val); |
| 95 | + } |
| 96 | + /** |
| 97 | + * @brief Returns comb across all values between l and r |
| 98 | + * @param l The left endpoint of the range |
| 99 | + * @param r The right endpoint of the range |
| 100 | + * @return The value of the comb operations |
| 101 | + */ |
| 102 | + T range_comb(int l, int r) { |
| 103 | + return range_comb(1, 1, size, l, r); |
| 104 | + } |
| 105 | +}; |
| 106 | +} // namespace data_structures |
| 107 | + |
| 108 | +/** |
| 109 | + * @brief Self-test implementations |
| 110 | + * @returns void |
| 111 | + */ |
| 112 | +static void test() { |
| 113 | + data_structures::SegmentTree<int> t(5); |
| 114 | + t.update(1, 1); |
| 115 | + t.update(2, 2); |
| 116 | + t.update(3, 3); |
| 117 | + t.update(4, 4); |
| 118 | + t.update(5, 5); |
| 119 | + assert(t.range_comb(1, 3) == 6); // 1 + 2 + 3 = 6 |
| 120 | + t.update(1, 3); |
| 121 | + assert(t.range_comb(1, 3) == 8); // 3 + 2 + 3 = 8 |
| 122 | + |
| 123 | + std::cout << "All tests have successfully passed!\n"; |
| 124 | +} |
| 125 | + |
| 126 | +/** |
| 127 | + * @brief Main function |
| 128 | + * @returns 0 on exit |
| 129 | + */ |
| 130 | +int main() { |
| 131 | + test(); // run self-test implementations |
| 132 | + return 0; |
| 133 | +} |
0 commit comments