Skip to content

Commit 170f7e5

Browse files
fgmnrealstealthninjaPanquesito7
authored
feat: Treap Data Structure (TheAlgorithms#2458)
* feat: Treap Data Structure * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: suggested changes * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * fix: add more self-tests * fix: suggested changes Co-authored-by: realstealthninja <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * suggested changes Co-authored-by: David Leal <[email protected]> * fix: add namespace prefixs * suggested changes Co-authored-by: David Leal <[email protected]> * fix: use containers * fix: add initialization * fix: suggested changes * fix: a few `clang-tidy` warnings * fix: segmentation fault --------- Co-authored-by: realstealthninja <[email protected]> Co-authored-by: David Leal <[email protected]>
1 parent 17a5326 commit 170f7e5

File tree

1 file changed

+259
-0
lines changed

1 file changed

+259
-0
lines changed

data_structures/treap.cpp

Lines changed: 259 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,259 @@
1+
/**
2+
* @file
3+
* @brief A balanced binary search tree (BST) on the basis of binary search tree
4+
* and heap: the [Treap](https://en.wikipedia.org/wiki/Treap) algorithm
5+
* implementation
6+
*
7+
* @details
8+
* Implementation of the treap data structre
9+
*
10+
* Support operations including insert, erase, and query (the rank of specified
11+
* element or the element ranked x) as the same as BST
12+
*
13+
* But these operations take O(log N) time, since treap keeps property of heap
14+
* using rotate operation, and the desired depth of the tree is O(log N).
15+
* There's very little chance that it will degenerate into a chain like BST
16+
*
17+
* @author [Kairao ZHENG](https://github.com/fgmn)
18+
*/
19+
20+
#include <array> /// For array
21+
#include <cassert> /// For assert
22+
#include <iostream> /// For IO operations
23+
24+
/**
25+
* @namespace
26+
* @brief Data Structures
27+
*/
28+
namespace data_structures {
29+
/**
30+
* @namespace
31+
* @brief Functions for the [Treap](https://en.wikipedia.org/wiki/Treap)
32+
* algorithm implementation
33+
*/
34+
namespace treap {
35+
const int maxNode = 1e5 + 5; ///< maximum number of nodes
36+
/**
37+
* @brief Struct representation of the treap
38+
*/
39+
struct Treap {
40+
int root = 0; ///< root of the treap
41+
int treapCnt = 0; ///< Total number of current nodes in the treap
42+
std::array<int, maxNode> key = {}; ///< Node identifier
43+
std::array<int, maxNode> priority = {}; ///< Random priority
44+
std::array<std::array<int, 2>, maxNode> childs = {
45+
{}}; ///< [i][0] represents the
46+
///< left child of node i, and
47+
///[i][1] represents the right
48+
std::array<int, maxNode> cnt =
49+
{}; ///< Maintains the subtree size for ranking query
50+
std::array<int, maxNode> size = {}; ///< The number of copies per node
51+
/**
52+
* @brief Initialization
53+
*/
54+
Treap() : treapCnt(1) {
55+
priority[0] = INT32_MAX;
56+
size[0] = 0;
57+
}
58+
/**
59+
* @brief Update the subtree size of the node
60+
* @param x The node to update
61+
*/
62+
void update(int x) {
63+
size[x] = size[childs[x][0]] + cnt[x] + size[childs[x][1]];
64+
}
65+
/**
66+
* @brief Rotate without breaking the property of BST
67+
* @param x The node to rotate
68+
* @param t 0 represent left hand, while 1 right hand
69+
*/
70+
void rotate(int &x, int t) {
71+
int y = childs[x][t];
72+
childs[x][t] = childs[y][1 - t];
73+
childs[y][1 - t] = x;
74+
// The rotation will only change itself and its son nodes
75+
update(x);
76+
update(y);
77+
x = y;
78+
}
79+
/**
80+
* @brief Insert a value into the specified subtree (internal method)
81+
* @param x Insert into the subtree of node x (Usually x=root)
82+
* @param k Key to insert
83+
*/
84+
void _insert(int &x, int k) {
85+
if (x) {
86+
if (key[x] == k) {
87+
cnt[x]++;
88+
} // If the node already exists, the number of copies is ++
89+
else {
90+
int t = (key[x] < k); // Insert according to BST properties
91+
_insert(childs[x][t], k);
92+
// After insertion, the heap properties are retained by rotation
93+
if (priority[childs[x][t]] < priority[x]) {
94+
rotate(x, t);
95+
}
96+
}
97+
} else { // Create a new node
98+
x = treapCnt++;
99+
key[x] = k;
100+
cnt[x] = 1;
101+
priority[x] = rand(); // Random priority
102+
childs[x][0] = childs[x][1] = 0;
103+
}
104+
update(x);
105+
}
106+
/**
107+
* @brief Erase a value from the specified subtree (internal method)
108+
* @param x Erase from the subtree of node x (Usually x=root)
109+
* @param k Key to erase
110+
*/
111+
void _erase(int &x, int k) {
112+
if (key[x] == k) {
113+
if (cnt[x] > 1) {
114+
cnt[x]--;
115+
} // If the node has more than one copy, the number of copies --
116+
else {
117+
if (childs[x][0] == 0 && childs[x][1] == 0) {
118+
x = 0;
119+
return;
120+
} // If there are no children, delete and return
121+
// Otherwise, we need to rotate the sons and delete them
122+
// recursively
123+
int t = (priority[childs[x][0]] > priority[childs[x][1]]);
124+
rotate(x, t);
125+
_erase(x, k);
126+
}
127+
} else { // Find the target value based on BST properties
128+
_erase(childs[x][key[x] < k], k);
129+
}
130+
update(x);
131+
}
132+
/**
133+
* @brief Find the KTH largest value (internal method)
134+
* @param x Query the subtree of node x (Usually x=root)
135+
* @param k The queried rank
136+
* @return The element ranked number k
137+
*/
138+
int _get_k_th(int &x, int k) {
139+
if (k <= size[childs[x][0]]) {
140+
return _get_k_th(childs[x][0], k);
141+
}
142+
k -= size[childs[x][0]] + cnt[x];
143+
if (k <= 0) {
144+
return key[x];
145+
}
146+
return _get_k_th(childs[x][1], k);
147+
}
148+
/**
149+
* @brief Query the rank of specified element (internal method)
150+
* @param x Query the subtree of node x (Usually x=root)
151+
* @param k The queried element
152+
* @return The rank of element k
153+
*/
154+
int _get_rank(int x, int k) {
155+
if (!x) {
156+
return 0;
157+
}
158+
if (k == key[x]) {
159+
return size[childs[x][0]] + 1;
160+
}
161+
else if (k < key[x]) {
162+
return _get_rank(childs[x][0], k);
163+
}
164+
else {
165+
return size[childs[x][0]] + cnt[x] + _get_rank(childs[x][1], k);
166+
}
167+
}
168+
/**
169+
* @brief Get the predecessor node of element k
170+
* @param k The queried element
171+
* @return The predecessor
172+
*/
173+
int get_predecessor(int k) {
174+
int x = root, pre = -1;
175+
while (x) {
176+
if (key[x] < k) {
177+
pre = key[x], x = childs[x][1];
178+
} else {
179+
x = childs[x][0];
180+
}
181+
}
182+
return pre;
183+
}
184+
/**
185+
* @brief Get the successor node of element k
186+
* @param k The queried element
187+
* @return The successor
188+
*/
189+
int get_next(int k) {
190+
int x = root, next = -1;
191+
while (x) {
192+
if (key[x] > k) {
193+
next = key[x], x = childs[x][0];
194+
} else {
195+
x = childs[x][1];
196+
}
197+
}
198+
return next;
199+
}
200+
/**
201+
* @brief Insert element (External method)
202+
* @param k Key to insert
203+
*/
204+
void insert(int k) { _insert(root, k); }
205+
/**
206+
* @brief Erase element (External method)
207+
* @param k Key to erase
208+
*/
209+
void erase(int k) { _erase(root, k); }
210+
/**
211+
* @brief Get the KTH largest value (External method)
212+
* @param k The queried rank
213+
* @return The element ranked number x
214+
*/
215+
int get_k_th(int k) { return _get_k_th(root, k); }
216+
/**
217+
* @brief Get the rank of specified element (External method)
218+
* @param k The queried element
219+
* @return The rank of element k
220+
*/
221+
int get_rank(int k) { return _get_rank(root, k); }
222+
};
223+
} // namespace treap
224+
} // namespace data_structures
225+
226+
/**
227+
* @brief Self-test implementations
228+
* @returns void
229+
*/
230+
static void test() {
231+
data_structures::treap::Treap mTreap; ///< Treap object instance
232+
233+
mTreap.insert(1);
234+
mTreap.insert(2);
235+
mTreap.insert(3);
236+
assert(mTreap.get_k_th(2) == 2);
237+
mTreap.insert(4);
238+
mTreap.insert(5);
239+
mTreap.insert(6);
240+
assert(mTreap.get_next(4) == 5);
241+
mTreap.insert(7);
242+
assert(mTreap.get_predecessor(7) == 6);
243+
mTreap.erase(4);
244+
assert(mTreap.get_k_th(4) == 5);
245+
assert(mTreap.get_rank(5) == 4);
246+
mTreap.insert(10);
247+
assert(mTreap.get_rank(10) == 7);
248+
assert(mTreap.get_predecessor(10) == 7);
249+
250+
std::cout << "All tests have successfully passed!\n";
251+
}
252+
/**
253+
* @brief Main function
254+
* @returns 0 on exit
255+
*/
256+
int main() {
257+
test(); // run self-test implementations
258+
return 0;
259+
}

0 commit comments

Comments
 (0)