# COSTEMP - Editorial

Author: Nishikant Parmar
Editorialist: Nishikant Parmar

EASY-MEDIUM

# PREREQUISITES:

Depth-First Search, Dynamic Programming on Trees

# PROBLEM:

Given a tree with N vertices numbered from 0 to N-1 and N-1 edges, output an array where the i^{th} element of the array denotes the sum of distances of the vertex i to all other vertices in the tree.

Note : Here the distance between two vertices connected through an edge is 1 units.

# QUICK EXPLANATION:

Let, nodesCount[v] be no. of nodes which are in the subtree of node v (including v) and ans[v] be the sum of distances of the node v to the remaining nodes.

Then, ans[v] = ans[u] + N - 2*nodesCount[v] (where, u is the parent of v)

# EXPLANATION:

Throughout the editorial consider u to be parent of v.

Suppose, down[v] denotes sum of distances to all nodes in the subtree of node v and let us define nodesCount[v], ans[v] same as above.

Calculation of nodesCount[u] :

nodesCount[u] = 1 + \sum\limits_{for\;all\;children\;v\;of\;u} nodesCount[v]

Calculation of down[u] :

down[u] = \sum\limits_{for\;all\;children\;v\;of\;u} (nodesCount[v] + down[v])

To get the sum of distances from u to all nodes in the subtree of v we need to add 1 unit extra distance for each node in subtree of v. Since, there are nodesCount[v] many nodes hence, we add this to down[v] to get the sum of distances from u to all nodes in subtree of v and take the sum over all children v of u.

Now, if we calculate up[v] i.e. sum of distances to all nodes from v that are not a part of subtree of v then,

ans[v] = down[v] + up[v]

Calculation of up[u] :

up[v] = ans[u] - (nodesCount[v] + down[v]) + (N - nodesCount[v])

Here, we assume that we already know ans[u].

When we subtract (nodesCount[v] + down[v]) from ans[u] we get sum of distances to all nodes from u to the nodes which are not in the subtree of v. Now, to make it the sum of distances from v we add the number of nodes which are not in subtree of v (This is similar to calculating down[u] from down[v] by adding nodesCount[v])

Calculation of ans[u]:

ans[v] = down[v] + up[v]

\implies ans[v] = down[v] + ans[u] - (nodesCount[v] + down[v]) + (N - nodesCount[v])

\implies ans[v] = ans[u] - N - 2*nodesCount[v]

We had assumed that we already know ans[u] while calculating up[v]. We can incrementally calculate ans[u] and use it for up[v] calculation using another Depth-First Search. Here, the base case is ans[root] = down[root], since for root all nodes lie in its subtree and up[root] = 0.

# SOLUTIONS:

Setter's Solution

#include <bits/stdc++.h>
using namespace std;

void DFS(int u, vector<vector<int>> &GRAPH, vector<bool> &visited, vector<int> &down, vector<int> &nodesCount){

visited[u] = true;
nodesCount[u]++;

for(int v:GRAPH[u]){
if(!visited[v]){
DFS(v, GRAPH, visited, down,nodesCount);
nodesCount[u]+=nodesCount[v];
down[u]+=nodesCount[v]+down[v];
}
}
}

void DFS2(int u, vector<vector<int>> &GRAPH, vector<bool> &visited,vector<int> &nodesCount, vector<int> &ans){

visited[u] = true;
int N = visited.size();

for(int v:GRAPH[u]){
if(!visited[v]){
ans[v] = ans[u] + N - 2*nodesCount[v];
DFS2(v,GRAPH,visited,nodesCount,ans);
}
}
}

int main(){

int N;
cin>>N;
vector<vector<int>> GRAPH(N);

for(int i=0;i<N-1;i++){
int u, v;
cin>>u;cin>>v;
GRAPH[u].push_back(v);
GRAPH[v].push_back(u);
}

// DFS 1
vector<int> down(N, 0);
vector<bool> visited(N, false);
vector<int> nodesCount(N, 0);
DFS(0, GRAPH, visited, down,nodesCount);

// DFS 2
vector<int> ans(N, 0);
for(int i=0;i<N;i++) visited[i] = false;
ans[0] = down[0];

DFS2(0,GRAPH,visited,nodesCount,ans);
for(int x:ans) cout << x << " ";

return 0;

}


Editorialist's Solution

import sys
sys.setrecursionlimit(10**6)

N = int(input())
GRAPH = [[] for i in range(N)]

for i in range(N-1):
u, v = map(int, input().split())
GRAPH[u].append(v)
GRAPH[v].append(u)

visited = [False for i in range(N)]
down = [0 for i in range(N)]
nodesCount = [0 for i in range(N)]

def DFS(u):
visited[u] = True
nodesCount[u]+=1

for v in GRAPH[u]:
if not visited[v]:
DFS(v)
nodesCount[u]+=nodesCount[v]
down[u]+=nodesCount[v]+down[v]

def DFS2(u):
visited[u] = True
for v in GRAPH[u]:
if not visited[v]:
ans[v] = ans[u] + N - 2*nodesCount[v]
DFS2(v)

DFS(0)

visited = [False for i in range(N)]
ans = [0 for i in range(N)]
ans[0] = down[0]
DFS2(0)

for x in ans:
print(x, end="")
print(" ", end="")

1 Like

Its moreover a question of BFS

this is a pretty straightforward solution.

import java.util.*;
import java.io.*;

class Cosmic_Temple_distance_codechef {

static int mod = (int) (1e9 + 7);
static boolean vis[];
static int dis[];
public static void main(String[] args) throws java.lang.Exception {
int t =1;
while (t-- > 0) {
for(int i=0; i<n; i++) adj[i]=new ArrayList<>();
vis=new boolean[n];dis=new int[n];
for(int i=0; i<n-1; i++) {
int u=Integer.parseInt(st.nextToken()),
v=Integer.parseInt(st.nextToken());
}
long sum;
for(int i=0; i<n; i++) {
Arrays.fill(vis, false);
Arrays.fill(dis, 0);
dfs(i,0);
sum=0;
for(int j=0; j<n; j++) {
sum+=dis[j];
}
System.out.print(sum+" ");
}
}
}
static void dfs(int v, int d) {
vis[v]=true;
dis[v]=d;