diff --git a/src/classes/tree/fenwick_tree.h b/src/classes/tree/fenwick_tree.h new file mode 100644 index 00000000..bf966899 --- /dev/null +++ b/src/classes/tree/fenwick_tree.h @@ -0,0 +1,63 @@ +#ifndef FENWICK_TREE_H +#define FENWICK_TREE_H + +#ifdef __cplusplus +#include +#include +#endif + +/** +* @brief fenwick tree class +*/ +template +struct fenwick_tree { + std::vector tree; + int n; + + /** + * @brief default constructor of fenwick tree class + * @param v: the input vector + */ + explicit fenwick_tree(const std::vector &v) noexcept : n(int(v.size())) { + tree = std::vector(n, 0); + for(int i = 0; iupdate(i, v[i]); + } + } + + /** + * @brief sum query function + * @param k: the ending index of the query + * @return T: the sum of range [0, k] + */ + T sum(int k) { + T sum = 0; + for(; k >= 0; k = (k & (k + 1)) - 1) { + sum += tree[k]; + } + return sum; + } + + /** + * @brief sum query function(from index a to b) + * @param a: starting index + * @param b: ending index + * @returns T: the sum of range [a, b] + */ + T sum(int a, int b) { + return sum(b) - sum(a - 1); + } + + /** + * @brief update query function + * @param k: the index + * @param x: the value that will be added to data[k] + */ + void update(int k, int x) { + for(; k < n; k = k | (k + 1)) { + tree[k] += x; + } + } +}; + +#endif diff --git a/tests/tree/fenwick_tree.cc b/tests/tree/fenwick_tree.cc new file mode 100644 index 00000000..0d21bd98 --- /dev/null +++ b/tests/tree/fenwick_tree.cc @@ -0,0 +1,37 @@ +#include "../../src/classes/tree/fenwick_tree.h" +#include "../../third_party/catch.hpp" + + +TEST_CASE("Testing fenwick tree default constructor") { + std::vector v = {1, 4, 5, 6, 7}; + CHECK_NOTHROW(fenwick_tree(v)); +} + +TEST_CASE("Testing fenwick tree sum query [1]") { + std::vector v = {1, 2, 3, 4, 5}; + fenwick_tree f(v); + REQUIRE(f.sum(2) == 6); + REQUIRE(f.sum(1) == 3); + REQUIRE(f.sum(4) == 15); + REQUIRE(f.sum(0) == 1); +} + +TEST_CASE("Testing fenwick tree sum query [2]") { + std::vector v = {1, 2, 3, 4, 5}; + fenwick_tree f(v); + REQUIRE(f.sum(2) == 6); + REQUIRE(f.sum(4) == 15); + REQUIRE(f.sum(1, 2) == 5); + REQUIRE(f.sum(3, 4) == 9); + REQUIRE(f.sum(1, 3) == 9); +} + +TEST_CASE("Testing update query [1]") { + std::vector v = {1, 2, 3, 4, 5}; + fenwick_tree f(v); + + f.update(0, 2); + REQUIRE(f.sum(0, 1) == 5); + f.update(0, -5); + REQUIRE(f.sum(0, 2) == 3); +} diff --git a/tutorial/tree/fenwick_tree.md b/tutorial/tree/fenwick_tree.md new file mode 100644 index 00000000..175161d3 --- /dev/null +++ b/tutorial/tree/fenwick_tree.md @@ -0,0 +1,25 @@ +### Mini tutorial for fenwick tree class + +-- fenwick_tree creates a fenwick tree with nodes of type T + +### *default constructor*: +```cpp +vector v = {1, 2, 3, 4, 5}; +fenwick_tree f(v); // this creates the fenwick tree + // with the passed vector v +``` + +### *sum query*: +```cpp +... +cout << f.sum(2) << '\n'; // this should return 6 +cout << f.sum(0, 2) << '\n'; // this should return 6 as well +``` + +### *update query*: +```cpp +... +f.update(0, -5); // updates the value of index 0 to -4(1 + -5) +cout << f.sum(2) << '\n'; // this should return 1 now +cout << f.sum(0, 2) << '\n'; // this should return 1 as well +```