COSTEMP - Editorial

editorial
, , , , ,
1

PROBLEM LINK:

Practice
Contest

Author: Nishikant Parmar
Editorialist: Nishikant Parmar

DIFFICULTY:

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
2

Its moreover a question of BFS

3

this is a pretty straightforward solution.

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

class Cosmic_Temple_distance_codechef {

	static int mod = (int) (1e9 + 7);
	static List<Integer> adj[];
	static boolean vis[];
	static int dis[];
	public static void main(String[] args) throws java.lang.Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int t =1;
		while (t-- > 0) {
			int n=Integer.parseInt(br.readLine());
			adj=new ArrayList[n+1];
			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++) {
				StringTokenizer st = new StringTokenizer(br.readLine());
				int u=Integer.parseInt(st.nextToken()),
				v=Integer.parseInt(st.nextToken());
				adj[u].add(v);adj[v].add(u);
			}
			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;
		for(int x: adj[v]) {
			if(!vis[x]) dfs(x, dis[v]+1);
		}
	}
}

you can check it also…

4

O(N^2) solution got AC.
Solution link: https://www.codechef.com/viewsolution/38320756

5

How BFS?

6

What is the logic behind this approach?