PROBLEM LINK:
Author: Sunny Aggarwal
Tester: Sergey Kulik
Editorialist: Sunny Aggarwal
DIFFICULTY:
Medium
PRE-REQUISITES:
Trees, Segment Trees, Implementation, Queries Handling, Data Structures.
PROBLEM STATEMENT:
Given a complete binary tree T consisting of N nodes numbered from 1 to N. For a node v, its left child is 2 \times v if it exists and its right child is 2 \times v + 1 if it exists. You are asked to process some type of queries on tree T as mentioned below.
- 1 X K Update the tree T such that node X becomes uncolored, all the nodes at distance 1 from X get colored with color 1, all nodes at distance 2 from X get colored with color 2, and so on up to nodes at distance K from X. Formally, color all nodes that are at distance (1 \le D \le K) units from node X with color D and make node X colourless.
Please note that when we say node at distance D from node X, you have to consider all the nodes in the tree, not just the ones in the subtree of X.
-
2 X Y Count the number of distinct colors on the unique simple path from node X to node Y.
-
3 X Count the number of distinct colors in the subtree rooted at node X.
SOLUTION:
Subtask 1:
This subtask is easy and scoring.
Create an adjacency list for the given tree and process the queries on the tree using any traversal algorithm like depth first search and breadth first search.
How to process queries of type 1 ?
Start from the given node U and color all the nodes that at most distance of K units from the given node U as shown below.
const int maxn = 1111;
vector adj[ maxn ];
int col[ maxn ], n;
void query1(int u, int k, int par = 0, int path_len = 0) {
// If distance > k units return
if( path_len > k ) {
return;
}
// color current node
col[u] = path_len;
for(auto it: adj[u]) {
if( it != par ) {
query1(it, k, u, path_len + 1);
}
}
}
Time Complexity: O(N)
How to handle queries to type 2 ?
We can simply run a traversal algorithm from one end of the given path to other end and can also maintain a list of colors on the path.
bool ok = false;
vector color_list;
void query2(int u, int par=0) {
if( u == v ) {
// if current node is other end
// of the path given in the query
ok = true;
}
for( auto it: adj[u] ) {
if( it != par && !ok ) {
query2(it, u);
}
}
if( ok ) {
color_list.push_back(col[u]);
}
}
Finding distinct colors in the given list is a trivial problem.
Time Complexity: O(N)
How to handle queries to type 2 efficiently?
We know that the given tree is complete binary tree and its height is atmost log_2{N}. This helps to claim that the distance between any pair of nodes (U, V) in the tree is atmost 2 \times \log_2{N}.
vector color_list;
void query2(int u, int v) {
while( u != v ) {
if(u < v) {
color_list.push_back(col[v]);
v /= 2;
} else {
color_list.push_back(col[u]);
u /= 2;
}
}
color_list.push_back(col[u]);
}
Time Complexity: O(log(N))
How to handle queries of type 3?
Suppose we need to find the number of distinct colors in the subtree rooted at node u which can be easily handled using the given procedure.
vector color_list;
void query3(int u, int par=u/2 ) {
color_list.push_back( col[u] );
for(auto it: adj[u]) {
if( it != par ) {
query3(it, u);
}
}
}
Note that the use of par=u/2 is neccessary otherwise it will explore the full tree T.
Subtask 2:
It should be noted that the procedure we used to handle queries of type 1 & 3 in the subtask 1 will surely time out for the subtask 2.
How to handle query 1 in this case ?
Let us try to update only a subtree rooted at a node U upto maximum distance of K units. Note that
- At 0^{th} level in subtree rooted at U, we have to update nodes having number [U \quad to \quad U] (i.e node U only) with color 0 ( 0 denotes colorless ).
- At 1^{st} level in subtree rooted at U, we have to update all nodes having number [2 \times U \quad to \quad 2 \times U + 1] with color 1.
- At 2^{nd} level in subtree rooted at U, we have to update all nodes having number [2 \times (2 \times U) \quad to \quad 2 \times (2 \times U + 1) + 1] with color 2.
- and so on upto min(K, number of levels in subtree rooted at U).
Note that as the given tree is a complete binary tree, its height will be atmost \log_2{N} and at each level we have a continuous segment from L to R to update with a single value and therefore, a segment tree along with lazy propagation can be used to perform updates in \log_2{N} time.
It should also be noted that in this problem query type 1 is not limited to subtree only so we also need to apply the above procedure for each ancestor of node U which adds another \log_2{N} factor to our update complexity.
To query efficiently, it is suggested to maintain a bitmask of 2 \times \log_2{N} bits (because this is maximum possible distance between any pair of nodes in the tree and hence the maximum color value) in each node of the segment tree such that operations like setting and unsetting a bit on a range can be performed efficiently.
You can learn about segment tree with lazy progagation here. Also try solving this problem.
How to handle queries to type 2?
Efficient method explained to handle queries of type 2 in first subtask is good enough.
How to handle queries to type 3?
The idea is almost similiar to what we have just explained while updating a subtree but in this case we are asked to query the number of distinct elements in the subtree which is equivalent of taking bitwise OR of all the elements in the subtree and then finding number of set bits in the resultant value. Note that bitwise OR for all the elements can be found by taking bitwise OR for all the segments in the subtree rooted at given node U.
Please check the setter’s solution for implementation details. Feel free to ask if anything is not clear.
TIME COMPLEXITY:
O((\log_2{N})^3 \times Q)
SETTER’S AND TESTER’S SOLUTION:
Setter’s solution can be found here
Tester’s solution can be found here