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

 Live Demo

#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, ]