CODEWITHSUNDEEP
c++cachinggraphdynamic-programming
Rohit gupta
Rohit gupta

Reputation: 281

How to optimize this graph/tree problem counting distance between two nodes in C++?

Problem is to find sum of distance between two nodes in a tree

INPUT: 6 [[0,1],[0,2],[2,3],[2,4],[2,5]]

shows nodes

enter image description here

OUTPUT: [8,12,6,10,10,10]

Explanation: The tree is shown above.

We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)

equals 1 + 1 + 2 + 2 + 2 = 8.

Hence, answer[0] = 8, and so on.

This is my code, I have solved it but this gives TLE and unable to optimize this solution

How I can optimize this graph problem.

class Solution {
public:
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        vector<int> g[n];
        map<pair<int,int>,int> m;

        for(auto i:edges){
            g[i[0]].push_back(i[1]);
            g[i[1]].push_back(i[0]);
        }

        for(int i=0;i<n;i++){ 
            queue<int> q;
            q.push(i);
            vector<bool> vis(n,0);
            vis[i]=1;
            int incr=1;
            while(!q.empty()){
                int k=q.size();
                while(k--){
                    int t=q.front();q.pop();
                    for(int j=0;j<g[t].size();j++){
                        if(!vis[g[t][j]] && g[t][j]!=i){
                            m[{i,g[t][j]}]+=incr;
                            q.push(g[t][j]);
                            vis[g[t][j]]=1;
                        }
                    }
                }
                incr++;
            }
        }
        vector<int> res(n,0);
        for(auto i:m){
            res[i.first.first]+=i.second;
        }
        return res;
    }
};

Upvotes: 0

Views: 191

Answers (1)

Shubham Sharma
Shubham Sharma

Reputation: 141

As I can see you are using bfs for every node to find the distance.
What you can do is use dynamic programming
Follow the steps below to solve the problem

  1. Initialize a vector dp to store the sum of the distances from each node i to all the leaf nodes of the tree.
  2. Initialize a vector leaves to store the count of the leaf nodes in the sub-tree of node i considering 1 as the root node.
  3. Find the sum of the distances from node i to all the leaf nodes in the sub-tree of i considering 1 as the root node using a modified Depth First Search Algorithm.
Let node a be the parent of node i

leaves[a] += leaves[i] ;
dp[a] += dp[i] + leaves[i]
  1. Use the re-rooting technique to find the distance of the remaining leaves of the tree that are not in the sub-tree of node i. To calculate these distances, use another modified Depth First Search (DFS) algorithm to find and add the sum of the distances of leaf nodes to node i.
Let a be the parent node and i be the child node, then
Let the number of leaf nodes outside the sub-tree i that are present in the sub-tree a be L

L = leaves[a] – leaves[i] ;
dp[i] += ( dp[a] – dp[i] ) + ( L – leaves[i] ) ;
leaves[i] += L ;

Upvotes: 2

Related Questions