diff --git a/DIRECTORY.md b/DIRECTORY.md index 681b0c3ba..605c4a1fe 100644 --- a/DIRECTORY.md +++ b/DIRECTORY.md @@ -204,6 +204,7 @@ ## Range Queries * [Fenwick Tree](https://github.com/TheAlgorithms/C-Plus-Plus/blob/master/range_queries/fenwick_tree.cpp) + * [Heavy Light Decomposition](https://github.com/TheAlgorithms/C-Plus-Plus/blob/master/range_queries/heavy_light_decomposition.cpp) * [Mo](https://github.com/TheAlgorithms/C-Plus-Plus/blob/master/range_queries/mo.cpp) * [Segtree](https://github.com/TheAlgorithms/C-Plus-Plus/blob/master/range_queries/segtree.cpp) diff --git a/range_queries/heavy_light_decomposition.cpp b/range_queries/heavy_light_decomposition.cpp new file mode 100644 index 000000000..b535a9808 --- /dev/null +++ b/range_queries/heavy_light_decomposition.cpp @@ -0,0 +1,639 @@ +/** + * @file + * @brief [Heavy Light + * Decomposition](https://en.wikipedia.org/wiki/Heavy_path_decomposition) + * implementation + * @author [Aniruthan R](https://github.com/aneee004) + * + * @details + * Heavy-Light Decomposition is a technique on trees, that supports the + * following: + * 1. Update node s, with a value v + * 2. Return the (sum) of all node values on the simple path from a to b + * (sum) can also be replced with XOR, OR, AND, min, or max + * + * The update is done in O(log n) time, and + * the query is done in O(log^2 n) time with HLD + * where, n is the number of nodes + * + * The template type is the data type of the value stored in the nodes. + * If a non-primitive data-type is used as a template, + * the coressponding operators must be overloaded. + * + * An HLD object can only be created with a constant number of nodes, and + * it cannot be changed later. Creaty an empty instance is not supported. + * + * To start answering updates and queries, + * 1. Create an instance of HLD object (obj), with the required data type. + * 2. Read in the edge/parent information and update it with obj.add_edge(). + * Note: The edges addes must be 0 indexed. + * 3. Create a vector with initial node values, and call set_node_val() with it. + * 4. Call obj.init() to populate the required information for supporting + * operations. + * 5. Call obj.update(node, new_val), to update the value at index 'node' to the + * new value. Note: node must be 0 indexed + * 6. Call obj.query(a, b) to get the (sum) of node values in the simple path + * from a to b. Note: a and b, must be 0 indexed. + * + * Sample I/O at the bottom. + * @todo Support edge weight queries, by storing the edge weight value in it's + * child algorithm verified by testing in CSES path queries: + * https://cses.fi/problemset/task/1138 + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/** + * @namespace range_queries + * @brief Algorithms and Data Structures that support range queries and updates. + */ +namespace range_queries { +/** + * @namespace heavy_light_decomposition + * @brief Heavy light decomposition algorithm + */ +namespace heavy_light_decomposition { +/** + * @brief A Basic Tree, which supports binary lifting + * @tparam the data type of the values stored in the tree nodes + * @details Deleting the default constructor + * An instance can only be created with the number of nodes + * Defaults: + * t_node indexing are zero based + * t_root is 0 + * depth of root_node is 0 + * Supports: + * lift :- lift a node k units up the tree + * kth_ancestor :- returns the kth ancestor + * lca :- returns the least common ancestor + */ +template class Tree { + // + +private: + std::vector> + t_adj; ///< an adjacency list to stores the tree edges + const int t_nodes, ///< number of nodes + t_maxlift; ///< maximum possible height of the tree + std::vector> + t_par; ///< a matrix to store every node's 2^kth parent + std::vector t_depth, ///< a vector to store the depth of a node, + t_size; ///< a vector to store the subtree size rooted at node + + int t_root; ///< the root of the tree + std::vector t_val; ///< values of nodes + template friend class HLD; + + /** + * @brief Utility function to compute sub-tree sizes + * @param u current dfs node + * @param p the parent of node @param u + * @returns void + */ + void dfs_size(int u, int p = -1) { + for (const int &v : t_adj[u]) { + if (v ^ p) { + dfs_size(v, u); + t_size[u] += t_size[v]; + } + } + } + + /** + * @brief Utility function to populate the t_par vector + * @param u current dfs node + * @param p the parent of node u + * @returns void + */ + void dfs_lca(int u, int p = -1) { + t_par[u][0] = p; + if (p != -1) { + t_depth[u] = 1 + t_depth[p]; + } + for (int k = 1; k < t_maxlift; k++) { + if (t_par[u][k - 1] != -1) { + t_par[u][k] = t_par[t_par[u][k - 1]][k - 1]; + } + } + + for (const int &v : t_adj[u]) { + if (v ^ p) { + dfs_lca(v, u); + } + } + } + +public: + /** + * @brief Class parameterized constructor, resizes the and initializes the + * data members + * @param nodes the total number of nodes in the tree + */ + explicit Tree(int nodes) + : t_nodes(nodes), t_maxlift(static_cast(floor(log2(nodes))) + 1) { + /* Initialize and resize all the vectors */ + t_root = 0; /* Default */ + t_adj.resize(t_nodes); + t_par.assign(t_nodes, std::vector(t_maxlift, -1)); + t_depth.assign(t_nodes, 0); + t_size.assign(t_nodes, 1); + t_val.resize(t_nodes); + } + + /** + * @brief Adds an undirected edge from node u to node v in the tree + * @param u the node where the edge is from + * @param v the node where the edge is to + * @returns void + */ + void add_edge(const int u, const int v) { + t_adj[u].push_back(v); + t_adj[v].push_back(u); + } + + /** + * @brief Set the root for the tree + * @param new_root the new root + * @returns void + */ + void change_root(int new_root) { t_root = new_root; } + + /** + * @brief Set the values for all the nodes + * @param node_val a vector of size n, with all the node values where, n is + * the number of nodes + * @returns void + */ + void set_node_val(const std::vector &node_val) { + assert(static_cast(node_val.size()) == t_nodes); + t_val = node_val; + } + + /** + * @brief This function must be called after the tree adjacency list and node + * values are populated The function initializes the required parameters, and + * populates the segment tree + * @returns void + */ + void init() { + assert(t_nodes > 0); + dfs_size(t_root); + dfs_lca(t_root); + } + + /** + * @brief The function lifts a node, k units up the tree. + * The lifting is done in place, and the result is stored in the address + * pointed by p. + * @param p a pointer to the variable that stores the node id + * @param dist the distance to move up the tree + * @returns void + */ + void lift(int *const p, int dist) { + for (int k = 0; k < t_maxlift; k++) { + if (*p == -1) { + return; + } + if (dist & 1) { + *p = t_par[*p][k]; + } + dist >>= 1; + } + } + + /** + * @brief The function returns the kth ancestor of a node + * @param p the node id whose kth ancestor is to be found + * @param dist the distance to move up the tree + * @returns the kth ancestor of node + */ + int kth_ancestor(int p, const int &dist) { + lift(&p, dist); + return p; + } + + /** + * @brief The function returns the least common ancestor of two nodes + * @param a node id_1 + * @param b node id_2 + * @returns the least common ancestor of node a, and node b + */ + int lca(int a, int b) { + assert(a >= 0 and b >= 0 and a < t_nodes and b < t_nodes); + if (t_depth[a] > t_depth[b]) { + lift(&a, t_depth[a] - t_depth[b]); + } + if (t_depth[b] > t_depth[a]) { + lift(&b, t_depth[b] - t_depth[a]); + } + if (a == b) { + return a; + } + for (int k = t_maxlift - 1; k >= 0; k--) { + if (t_par[a][k] != t_par[b][k]) { + a = t_par[a][k]; + b = t_par[b][k]; + } + } + return t_par[a][0]; + } +}; + +/** + * @brief Segment Tree, to store heavy chains + * @tparam the data type of the values stored in the tree nodes + */ +template class SG { +private: + /** + * @brief Everything here is private, + * and can only be accessed through the methods, + * in the derived class (HLD) + */ + + std::vector s_tree; ///< the segment tree, stored as a vector + int s_size; ///< number of leaves in the segment tree + X sret_init = 0; ///< inital query return value + template friend class HLD; + + /** + * @brief Function that specifies the type of operation involved when segments + * are combined + * @param lhs the left segment + * @param rhs the right segment + * @returns the combined result + */ + X combine(X lhs, X rhs) { return lhs + rhs; } + + /** + * @brief Class parameterized constructor. Resizes the and initilizes the data + * members. + * @param nodes the total number of nodes in the tree + * @returns void + */ + explicit SG(int size) { + s_size = size; + s_tree.assign(2 * s_size, 0ll); + } + + /** + * @brief Update the value at a node + * @param p the node to be udpated + * @param v the update value + * @returns void + */ + void update(int p, X v) { + for (p += s_size; p > 0; p >>= 1) { + s_tree[p] += v; + } + } + + /** + * @brief Make a range query from node label l to node label r + * @param l node label where the path starts + * @param r node label where the path ends + * @returns void + */ + X query(int l, int r) { + X lhs = sret_init, rhs = sret_init; + for (l += s_size, r += s_size + 1; l < r; l >>= 1, r >>= 1) { + if (l & 1) { + lhs = combine(lhs, s_tree[l++]); + } + if (r & 1) { + rhs = combine(s_tree[--r], rhs); + } + } + return combine(lhs, rhs); + } + + /** + * @brief Set the initialization for the query data type, based on requirement + * + * @details + * Change the sret_init, based on requirement: + * * Sum Query: 0 (Default) + * * XOR Query: 0 (Default) + * * Min Query: Infinity + * * Max Query: -Infinity + * @param new_sret_init the new init + */ + void set_sret_init(X new_sret_init) { sret_init = new_sret_init; } +}; + +/** + * @brief The Heavy-Light Decomposition class + * @tparam the data type of the values stored in the tree nodes + */ +template class HLD : public Tree, public SG { +private: + int label; ///< utility member to assign labels in dfs_labels() + std::vector h_label, ///< stores the label of a node + h_heavychlid, ///< stores the heavy child of a node + h_parent; ///< stores the top of the heavy chain from a node + + /** + * @brief Utility function to assign heavy child to each node (-1 for a leaf + * node) + * @param u current dfs node + * @param p the parent of node u + * @returns void + */ + void dfs_hc(int u, int p = -1) { + int hc_size = -1, hc_id = -1; + for (const int &v : Tree::t_adj[u]) { + if (v ^ p) { + dfs_hc(v, u); + if (Tree::t_size[v] > hc_size) { + hc_size = Tree::t_size[v]; + hc_id = v; + } + } + } + h_heavychlid[u] = hc_id; + } + + /** + * @brief Utility function to assign highest parent that can be reached though + * heavy chains + * @param u current dfs node + * @param p the parent of node u + * @returns void + */ + void dfs_par(int u, int p = -1) { + if (h_heavychlid[u] != -1) { + h_parent[h_heavychlid[u]] = h_parent[u]; + dfs_par(h_heavychlid[u], u); + } + for (const int &v : Tree::t_adj[u]) { + if (v ^ p and v ^ h_heavychlid[u]) { + dfs_par(v, u); + } + } + } + + /** + * @brief Utility function to lable the nodes so that heavy chains have a + * contigous lable + * @param u current dfs node + * @param p the parent of node u + * @returns void + */ + void dfs_labels(int u, int p = -1) { + h_label[u] = label++; + if (h_heavychlid[u] != -1) { + dfs_labels(h_heavychlid[u], u); + } + for (const int &v : Tree::t_adj[u]) { + if (v ^ p and v ^ h_heavychlid[u]) { + dfs_labels(v, u); + } + } + } + + /** + * @brief Utility function to break down a path query into two chain queries + * @param a node where the path starts + * @param b node where the path ends + * a and b must belong to a single root to leaf chain + * @returns the sum of ndoe values in the simple path from a to b + */ + X chain_query(int a, int b) { + X ret = SG::sret_init; + if (Tree::t_depth[a] < Tree::t_depth[b]) { + std::swap(a, b); + } + while (Tree::t_depth[a] >= Tree::t_depth[b]) { + int l = h_label[h_parent[a]]; + int r = h_label[a]; + if (Tree::t_depth[h_parent[a]] < Tree::t_depth[b]) { + l += Tree::t_depth[b] - Tree::t_depth[h_parent[a]]; + } + ret = SG::combine(ret, SG::query(l, r)); + a = Tree::t_par[h_parent[a]][0]; + if (a == -1) { + break; + } + } + return ret; + } + +public: + /** + * @brief Class parameterized constructor. Resizes the and initilizes the data + * members. + * @param nodes the total number of nodes in the tree + */ + explicit HLD(int nodes) : Tree(nodes), SG(nodes) { + /* Initialization and resize vectors */ + label = 0; + h_label.assign(Tree::t_nodes, -1); + h_heavychlid.assign(Tree::t_nodes, -1); + h_parent.resize(Tree::t_nodes); + iota(h_parent.begin(), h_parent.end(), 0); + } + + /** + * @brief This function must be called after the tree adjacency list and node + * values are populated The function initializes the required parametes, and + * populates the segment tree + * @returns void + */ + void init() { + Tree::init(); + + // Fill the heavy child, greatest parent, and labels + label = 0; + dfs_hc(Tree::t_root); + dfs_par(Tree::t_root); + dfs_labels(Tree::t_root); + + // Segment Tree Initialization + for (int i = 0; i < Tree::t_nodes; i++) { + SG::s_tree[h_label[i] + Tree::t_nodes] = Tree::t_val[i]; + } + for (int i = Tree::t_nodes - 1; i > 0; i--) { + SG::s_tree[i] = + SG::combine(SG::s_tree[i << 1], SG::s_tree[i << 1 | 1]); + } + } + + /** + * @brief This function updates the value at node with val + * @param node the node where the update is done + * @param val the value that is being updated + * @returns void + */ + void update(int node, X val) { + X diff = val - Tree::t_val[node]; + SG::update(h_label[node], diff); + Tree::t_val[node] = val; + } + + /** + * @brief This function returns the sum of node values in the simple path from + * from node_1 to node_2 + * @param a the node where the simple path starts + * @param b the node where the simple path ends + * (parameters are interchangeable, i.e., the function is commutative) + * @returns the sum of node values in the simple path from a to b + */ + X query(int a, int b) { + int lc = Tree::lca(a, b); + X ret = SG::sret_init; + assert(lc != -1); + ret += chain_query(a, lc); + ret += chain_query(b, lc); + return ret - Tree::t_val[lc]; + } +}; +} // namespace heavy_light_decomposition +} // namespace range_queries + +/** + * Test implementations + * @returns none + */ +static void test_1() { + std::cout << "Test 1:\n"; + + // Test details + int n = 5; + std::vector node_values = {4, 2, 5, 2, 1}; + std::vector> edges = {{1, 2}, {1, 3}, {3, 4}, {3, 5}}; + std::vector> queries = { + {2, 1, 4}, + {1, 3, 2}, + {2, 1, 4}, + }; + std::vector expected_result = {11, 8}; + std::vector code_result; + + range_queries::heavy_light_decomposition::HLD hld(n); + hld.set_node_val(node_values); + for (int i = 0; i < n - 1; i++) { + int u = edges[i][0], v = edges[i][1]; + hld.add_edge(u - 1, v - 1); + } + hld.init(); + for (const auto &q : queries) { + int type = q[0]; + if (type == 1) { + int p = q[1], x = q[2]; + hld.update(p - 1, x); + } else if (type == 2) { + int a = q[1], b = q[2]; + code_result.push_back(hld.query(a - 1, b - 1)); + } else { + continue; + } + } + for (int i = 0; i < static_cast(expected_result.size()); i++) { + assert(expected_result[i] == code_result[i]); + } + std::cout << "\nTest 1 passed!\n"; +} + +/** + * Second test implementations + * @returns void + */ +static void test_2() { + std::cout << "Test 2:\n"; + + // Test details (Bamboo) + int n = 10; + std::vector node_values = {1, 8, 6, 8, 6, 2, 9, 2, 3, 2}; + std::vector> edges = { + {10, 5}, {6, 2}, {10, 7}, {5, 2}, {3, 9}, {8, 3}, {1, 4}, {6, 4}, {8, 7}}; + std::vector> queries = { + {2, 1, 10}, {2, 1, 6}, {1, 3, 4}, {2, 1, 9}, {1, 5, 3}, + {1, 7, 8}, {2, 1, 4}, {2, 1, 8}, {1, 1, 4}, {1, 2, 7}}; + std::vector expected_result = {27, 11, 45, 9, 34}; + std::vector code_result; + + range_queries::heavy_light_decomposition::HLD hld(n); + hld.set_node_val(node_values); + for (int i = 0; i < n - 1; i++) { + int u = edges[i][0], v = edges[i][1]; + hld.add_edge(u - 1, v - 1); + } + hld.init(); + for (const auto &q : queries) { + int type = q[0]; + if (type == 1) { + int p = q[1], x = q[2]; + hld.update(p - 1, x); + } else if (type == 2) { + int a = q[1], b = q[2]; + code_result.push_back(hld.query(a - 1, b - 1)); + } else { + continue; + } + } + for (int i = 0; i < static_cast(expected_result.size()); i++) { + assert(expected_result[i] == code_result[i]); + } + std::cout << "\nTest2 passed!\n"; +} + +/** + * Third test implementations + * @returns void + */ +static void test_3() { + std::cout << "Test 3:\n"; + + // Test details + int n = 8; + std::vector node_values = {1, 8, 6, 8, 6, 2, 9, 2}; + std::vector> edges = {{1, 2}, {2, 3}, {3, 4}, {1, 5}, + {6, 3}, {7, 5}, {8, 7}}; + std::vector> queries = { + {2, 6, 8}, {2, 3, 6}, {1, 3, 4}, {2, 7, 1}, {1, 5, 3}, + {1, 7, 8}, {2, 6, 4}, {2, 7, 8}, {1, 1, 4}, {1, 2, 7}}; + std::vector expected_result = {34, 8, 16, 14, 10}; + std::vector code_result; + + range_queries::heavy_light_decomposition::HLD hld(n); + hld.set_node_val(node_values); + for (int i = 0; i < n - 1; i++) { + int u = edges[i][0], v = edges[i][1]; + hld.add_edge(u - 1, v - 1); + } + hld.init(); + for (const auto &q : queries) { + int type = q[0]; + if (type == 1) { + int p = q[1], x = q[2]; + hld.update(p - 1, x); + } else if (type == 2) { + int a = q[1], b = q[2]; + code_result.push_back(hld.query(a - 1, b - 1)); + } else { + continue; + } + } + for (int i = 0; i < static_cast(expected_result.size()); i++) { + assert(expected_result[i] == code_result[i]); + } + std::cout << "\nTest3 passed!\n"; +} + +/** + * Main function + */ +int main() { + test_1(); + test_2(); + test_3(); + return 0; +}