CONTEST LINK
Author: srmalkan
Editorialist:srmalkan
PREREQUISITES
Least common anscestor, LCA-Binary lifting
PROBLEM
Given a tree with n nodes and n-1 edges,for each query find number of edges between two nodes.
Explanation -
NAIVE approach:
->Implement BFS/DFS for each query by rooting one of the node.
->why dfs works?
–>because there exists only one path between any two nodes.
->why this approach don’t work?
–>let’s evaluate time complexity
→ time complexity of bfs/dfs : O(n)
→ Q can range from 1 to 10^6 : O(q)
→ total time complexity O(q*n)
–>worst case : O(10^6 * 10^5) → O(10^11)
–>The time limit for the problem is 1sec (Maximum 10^8 computations)
–>Since 10^11 >> 10^8, apprach fail to pass all the test cases.
OPTIMAL approach:
->Implement LCA(binary lifting) by rooting any node
->Ancestor is the node which you can reach by following parent up the tree to get to a position
→ Naive
–>store the depth and parent of each node
–>to find LCA between a and b
while a is not equal to b:
if(depth[a]>depth[b]):
a=parent[a]
else if depth[b]>depth[a]:
b=parent[b]
else:
a=parent[a]
b=parent[b]
–>O(diff in depth)
→ Binary Lifting
–>store the parents at height of powers of 2
–>using equation
parent[i][d] = parent[parent[i][d-1]][d-1]
–>psuedo code: memotization
// N is number of nodes
// D maximum depth possible ceil(log2(N))
par[N][D] = {-1}
depth=[0,0,...0] //N
visited=[false,false,.....false] //N
function dfs(int v):
visited[v] = true;
for d=1 to D:
parent = par[v][d-1];
if(parent!=-1): // to make sure we dont over step
par[v][d] = par[parent][d-1]
for u: adj[v]:
if not visited:
depth[u] = depth[v]+1; // precomputing depth of each node
par[u][0] = v; // precomuting ancestor for node u at height 2^"0" = 1 ...i.e its immediate parent
dfs(u)
–>psuedo-code : LCA
function lca(u,v):
if(depth[u]>depth[v]):
swap(u,v)
for d=D to 0 :
if level[v]-pow(2,d) >= level[u]: // bring efficiently both nodes at equal level
v=par[v][d]
if u == v:
return v
for d=D to 0:
if(par[u][d]!=par[v][d]): // make highest jump possible such that two nodes are not equal..such that we end up just LCA.
u = par[u][d]
v = par[v][d]
return par[u][0] // immediate parent of computed node is LCA;
→ Time complexity with q queries: O(log(n)*q)
SOLUTION
CLICK ME
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
ll depth[100100]={0};
vector<vector<ll>> adj(100100);
ll N,D;
vector<vector<ll>> par(100100,vector<ll>(50)); //initialize to -1;
void ingraph(int N,int M){
for(ll i=0;i<M;i++){
ll u,v;
cin>>u>>v;
adj[u-1].push_back(v-1);
adj[v-1].push_back(u-1);
}
}
void dfs(ll v){
for(ll d=1;d<=D;d++){
ll parent = par[v][d-1];
if(parent!=-1){
if(par[parent][d-1]!=-1){
par[v][d] = par[parent][d-1];
}
}
}
for(auto u : adj[v]){
if(u!=par[v][0]){
depth[u] = depth[v]+1;
par[u][0] = v;
dfs(u);
}
}
}
ll lca(ll u,ll v){
if(depth[u]>depth[v]){
ll temp = u;
u=v;
v=temp;
}
for(ll d=D;d>=0;d--){
if(depth[v]-pow(2,d) >= depth[u]){
v=par[v][d];
}
}
if(u==v){
return u;
}
for(ll d=D;d>=0;d--){
if(par[u][d]!=par[v][d]){
u=par[u][d];
v=par[v][d];
}
}
return par[u][0];
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0); cout.tie(0);
ll T;
cin>>T;
while(T--){
ll Q;
cin>>N>>Q;
D=ceil(log2(N));
for(ll i=0;i<N;i++){
adj[i].clear();
for(ll d=0;d<=D;d++){
par[i][d]=-1;
}
}
ingraph(N,N-1);
depth[0]=0;
dfs(0);
while(Q--){
ll a,b,d;
cin>>a>>b>>d;
ll x=lca(a-1,b-1);
if(depth[a-1]+depth[b-1]-2*depth[x]<=d){
cout<<"YES"<<"\n";
}else{
cout<<"NO"<<"\n";
}
}
}
return 0;
}
If you have any doubt, comment below in this thread.