题目链接:https://leetcode.com/problems/sum-of-distances-in-tree/
题意:给出一棵树,对于每个节点,求出其他所有节点与该节点的距离之和。节点数不超过10000。
思路:如图所示,暴力复杂度太高,考虑进行优化,如图所示,考虑一对父子节点u与v,son[u]表示以u为根节点的子树的节点个数。若已知父节点u到其他所有节点的距离和,则可推出子节点到其他所有节点的距离和:以子节点为根的树上的点到该子节点的距离是到其父节点的距离-1,而其他节点到该子节点的距离是到其父节点的距离+1(图中黄色和红色部分)。即$ans[v]= ans[u]-son[v]+cnt-son[v]$,其中$cnt$为节点总数。每个节点的子节点数可通过一次dfs求出,求所有节点的距离和也可以一次dfs求出:先求父节点,再求子节点。
class Solution { public: //dfs求出每个节点的子节点(包括自己) void dfs(int pre,int u, vector<int> g[],int height){ int s0=0; son[u]=1; for(int i=0;i<g[u].size();i++){ int v = g[u][i]; if(v==pre) continue; ans[0] +=height+1; dfs(u,v,g,height+1); son[u]+=son[v]; } } //dfs求出每个节点到其他所有节点的距离和 void Dfs(int pre,int u,vector<int> g[]){ for(int i=0;i<g[u].size();i++){ int v = g[u][i]; if(v==pre) continue; ans[v]= ans[u]-son[v]+cnt-son[v]; Dfs(u,v,g); } } vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { //总节点数 cnt = N; vector<int> g[N]; for(int i=0;i<edges.size();i++){ g[edges[i][0]].push_back(edges[i][1]); g[edges[i][1]].push_back(edges[i][0]); } vector <int> s; if(N==1){ s.push_back(0); return s; } son = new int[N]; ans = new int[N]; memset(son, 0, 4*N ); ans[0]=0; dfs(-1,0, g ,0); s.push_back(ans[0]); Dfs(-1,0,g); for(int i=1;i<N;i++) s.push_back(ans[i]); return s; } private: int *son = nullptr,*ans; int cnt; };