Data Structure
Networking
RDBMS
Operating System
Java
MS Excel
iOS
HTML
CSS
Android
Python
C Programming
C++
C#
MongoDB
MySQL
Javascript
PHP
Physics
Chemistry
Biology
Mathematics
English
Economics
Psychology
Social Studies
Fashion Studies
Legal Studies
- Selected Reading
- UPSC IAS Exams Notes
- Developer's Best Practices
- Questions and Answers
- Effective Resume Writing
- HR Interview Questions
- Computer Glossary
- Who is Who
Sum of Distances in Tree in C++
Suppose we have one undirected, connected tree where N nodes are present. These are labelled as 0...N-1 and N-1 edges are given. The ith edge connects nodes edges[i][0] and edges[i][1] together. We have to find a list where ans[i] is the sum of the distances between node i and all other nodes.
So, if the input is like N = 6 and edges = [(0,1),(0,2),(2,3),(2,4),(2,5)], then the output will be [8,12,6,10,10,10]
To solve this, we will follow these steps −
Define a function dfs1(), this will take node, parent,
for initialize i := 0, when i < size of graph[node], update (increase i by 1), do −
child := graph[node, i]
if child is not equal to parent, then −
dfs1(child, node)
cnt[node] := cnt[node] + cnt[child]
ans[node] := ans[node] + cnt[child] + ans[child]
Define a function dfs2(), this will take node, parent,
for initialize i := 0, when i < size of graph[node], update (increase i by 1), do−
child := graph[node, i]
if child is not equal to parent, then −
ans[child] := ans[node] - cnt[child] + N - cnt[child]
dfs2(child, node
Define an array ans
Define an array cnt
Define an array graph with 10005 rows
From the main method, do the following −
N of this := N
ans := Define an array of size N
cnt := Define an array of size N, fill this with 1
n := size of edges
for initialize i := 0, when i < n, update (increase i by 1), do −
u := edges[i, 0]
v := edges[i, 1]
insert v at the end of graph[u]
insert u at the end of graph[v]
dfs1(0, -1)
dfs2(0, -1)
return ans
Let us see the following implementation to get better understanding −
Example
#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
cout << "[";
for(int i = 0;
i<v.size(); i++){
cout << v[i] << ", ";
}
cout << "]"<<endl;
}
class Solution {
public:
void dfs1(int node, int parent) {
for (int i = 0; i < graph[node].size(); i++) {
int child = graph[node][i];
if (child != parent) {
dfs1(child, node);
cnt[node] += cnt[child];
ans[node] += cnt[child] + ans[child];
}
}
}
void dfs2(int node, int parent) {
for (int i = 0; i < graph[node].size(); i++) {
int child = graph[node][i];
if (child != parent) {
ans[child] = ans[node] - cnt[child] + N - cnt[child];
dfs2(child, node);
}
}
}
vector<int> ans;
vector<int> cnt;
vector<int> graph[10005];
int N;
vector<int> sumOfDistancesInTree(int N, vector<vector<int> >& edges) {
this->N = N;
ans = vector<int>(N);
cnt = vector<int>(N, 1);
int n = edges.size();
for (int i = 0; i < n; i++) {
int u = edges[i][0];
int v = edges[i][1];
graph[u].push_back(v);
graph[v].push_back(u);
}
dfs1(0, -1);
dfs2(0, -1);
return ans;
}
};
main(){
Solution ob;
vector<vector<int>> v = {{0,1},{0,2},{2,3},{2,4},{2,5}}; print_vector(ob.sumOfDistancesInTree(6, v));
}Input
{{0,1},{0,2},{2,3},{2,4},{2,5}}Output
[8, 12, 6, 10, 10, 10, ]
604 Views
