Skip to content

Commit 1a1570d

Browse files
feat: Segment Tree Data Structure (#2444)
* Segment Tree * fix: conformed to guidelines * fix: changed int to template type T in few places * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fixed: suggested changes Co-authored-by: David Leal <[email protected]> * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fix: suggested changes Co-authored-by: David Leal <[email protected]> * fix: remove duplicate line * Update segment_tree.cpp * chore: apply suggestions from code review --------- Co-authored-by: David Leal <[email protected]>
1 parent c0c2715 commit 1a1570d

File tree

1 file changed

+133
-0
lines changed

1 file changed

+133
-0
lines changed

data_structures/segment_tree.cpp

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
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

Comments
 (0)