PROBLEM LINK :
Setter: Abhilash
Tester: Harsh Raj
Editorialist: Harsh Raj
DIFFICULTY:
Easy-Medium
PREREQUISITES:
Basic graph theory and DFS Traversal
EXPLAINATION:
We answer the queries offline.
Do a DFS to label all the nodes with their depths from node.
For each node in query, store its corresponding depth value to
be queried. Do a second DFS from root, when you visit a node
that has to be queried, find the depths of nodes that can be
solution to this query and store in some global lookup table.
In given example, when we are at node 1, we update lookup
table of 1 with 1, i.e. all the nodes in subtree with depth
1 will be answer to query 1.
If the current depth is answer to some query, update answer
to the specific query.
when we reach node 2, we see its depth is 1, which is answer
to query 1, so we increase the count of query 1, by 1.
When returning back from subnodes (in
dfs) remove the added depth query from the global lookup
table.
Finally output all answers.
TIME COMPLEXITY:
O(N) as It is a tree and we a doing a DFS.
Settler's Solution
#include<bits/stdc++.h>
using namespace std;
vector<int>edges[100005];
vector<int>level_nodes[100005];
vector<int>start(100005),finish(100005),level(100005);
int counter=0;
void dfs(int s,int parent=-1,int _level=0)
{
counter++;
start[s]=counter;
level_nodes[_level].push_back(counter);
level[s]=_level;
for(auto x:edges[s])
{
if(x!=parent)
dfs(x,s,_level+1);
}
finish[s]=counter;
}
main()
{
int n,q;
cin>>n>>q;
for(int i=1;i<=n-1;i++){
int u,v;
cin>>u>>v;
edges[u].push_back(v);
edges[v].push_back(u);
}
dfs(1);
for(int i=0;i<=n;i++)
sort(level_nodes[i].begin(),level_nodes[i].end());
while(q--)
{
int node,depth;
cin>>node>>depth;
int _level=level[node]+depth;
if(_level>n)
{
cout<<0<<"\n";
continue;
}
auto r=upper_bound(level_nodes[_level].begin(),level_nodes[_level].end(),finish[node]);
auto l=lower_bound(level_nodes[_level].begin(),level_nodes[_level].end(),start[node]);
cout<<r-l<<"\n";
}
}
Tester's Solution
#include<bits/stdc++.h>
using namespace std;
#define ll long long int
#define eb emplace_back
const int N=2e5+4;
int n,q,depth[N],ans_query[N],query[N];
vector<int> g[N],down_query[N];
set<int> up_query[N];
// answer queries offline using single dfs
// down_query stores idexes of queries from the given index(source)
// query stores depth from the source of the given query index
// up_query stores indexes of queries for which the current node is one of the answer
void dfs1(int node,int par,int dep);
void dfs2(int node,int par);
int main(){
cin>>n>>q;
for(int i=1,to,from;i<n;i++){//store graph
cin>>to>>from;
g[to].eb(from),g[from].eb(to);
}
dfs1(1,0,0); //mark depths of all nodes
for(int i=1,source,dep;i<=q;i++){//store queries
cin>>source>>dep;
down_query[source].eb(i);
query[i]=dep;
}
dfs2(1,0);//find answer to queries
for(int i=1;i<=q;i++)
cout<<ans_query[i]<<"\n";
return 0;
}
void dfs2(int node,int par){
for(auto it:up_query[depth[node]]){//all the query indexes to which the current node is answer
ans_query[it]++;
}
for(auto it:down_query[node]){//contains query number
int dep=query[it]; //depth from given source
dep+=depth[node];//depth from root
up_query[dep].insert(it);//in the current subtree,at given depth,the node is answer to it'th query
}
for(auto it:g[node]){
if(it==par)
continue;
dfs2(it,node);
}
for(auto it:down_query[node]){//contains query number
int dep=query[it];
dep+=depth[node];
up_query[dep].erase(it);//it'th query is complete as we are not allowed to go towards root and all subnodes have been visited
}
}
void dfs1(int node,int par,int dep){
for(auto it:g[node]){
if(it==par)
continue;
dfs1(it,node,dep+1);
}
depth[node]=dep;//store depth of all nodes
}