Doubts in editorial of SPTREE using LCA

I have some doubts in the first dfs function that is called to root the tree at ‘a’ and to mark nodes having special node in their subtree :

EDIT : This is the same code used in the video editorial available on the official codechef channel

	public int dfs(int v , int depth , int parent, boolean isSpec[] , int par[] , int dep[] , int ans[] , int node[],ArrayList<Integer> graph[]){
	    dep[v] = depth;
	    par[v] = parent;
	    int sp = -1;
	    for(int u : graph[v]){
	            int temp = dfs(u,depth+1,v,isSpec,par,dep,ans,node,graph);
	            if(temp!=-1&&sp==-1) sp = temp; // why are we not reassigning sp to temp when sp is not -1 , why cant we use - if(temp!=-1) sp = temp here 
	    if(isSpec[v]) sp = v; // when we dont have a special node in subtree , then it makes sense to assign sp to v if v is a special node , but if we do have a special node in subtree then why are we reassigning sp to v if v is a special node ? 
	        ans[v]= dep[v];
	        node[v] = sp;
	    return sp;


  1. if(temp!=-1)sp=temp; This also works, no issues.
  2. Your second point is also valid, but I would like to state that node v is also part of subtree of v.


one last doubt ,

dfs(a); //Root the tree at 'a'
        for (int i = 0; i < n; i++) {
            if (ans[i] == -1 and par[i] != -1 and node[par[i]] != -1) {
                dfs2(i, 1, par[i], par[i]);

here in the if condition can we only use - if(ans[i]==-1) and proceed

EDIT : getting a tle with only if(ans[i]==-1) , and I think its because if we use only if(ans[i]==-1) then we are pairing nodes up randomly which leads to tle , whereas if we use other two checks we are surely gonna use node having special node in subtree as parent which is what we want to do

Basically you run the dfs2 for subtree of nodes that dont have any special node (ans[i]==-1). You need to pair all nodes in such subtree with closest ancestor, say L having special node in its subtree. Hence we run the dfs2 on i, such that ans[i]==-1 as well as parent of i exists and is L.

1 Like

Hey can you please help me out . I have replicated the exact same implementation as you did but in py3.6. I am getting a runtime error.

Here is my code

cook your dish here

Tc = int(input())

def dfs(v, depth = 0, parent = -1):
dep[v] = depth
par[v] = parent
sp = -1
for u in adj[v]:
if (u != par[v]):
temp = dfs(u, depth+1, v)
if (temp != -1):
sp = temp

if (is_spec[v]):
    sp = v
if (sp != -1):
    ans[v] = dep[v]
    node[v] = sp

return sp

def dfs2(v, depth, L, parent):
ans[v] = ans[L] - depth
node[v] = node[L]
for u in adj[v]:
if (u != parent):
dfs2(u, depth+1, L, v)

while (Tc > 0):
Tc = Tc - 1
params = input().split()
N = int(params[0]) #no of nodes
K = int(params[1]) #no of special nodes
a = int(params[2]) #node used in the formula
a -= 1
k = [int(x)-1 for x in input().split()] #list of k special nodes

is_spec = [False] * N
par = [-1] * N
ans = [-1] * N
adj = {} #dictionary to for adjacent neighbors
dep = [-1] * N
node = [-1] * N

for i in range(K): #take special nodes
    is_spec[k[i]] = True

for i in range(N-1):
    uv = input().split()
    u = int(uv[0]) - 1
    v = int(uv[1]) - 1
    if u not in adj.keys():
        adj[u] = []
    if v not in adj.keys():
        adj[v] = []
for i in range(N):
    if (ans[i] == -1) and (par[i] != -1) and (node[par[i]] != -1):
        dfs2(i, 1, par[i], par[i])
for i in range(N):
    if (i == N-1):
        print(ans[i], end="\n")
        print(ans[i], end=" ")

for i in range(N):
    if (i == N-1):
        print(node[i]+1, end="")
        print(node[i]+1, end=" ")

Kindly share your Codechef submission link.

Let me know if you are able to access this link. Thank you so much for helping out.