本文共 4132 字，大约阅读时间需要 13 分钟。
Leetcode 834. 树中距离之和
题目
给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
示例:
输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10] 解释: 如下为给定的树的示意图： 0 / \ 1 2 /|\ 3 4 5我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。说明: 1 <= N <= 10000
设dp[cur]表示当前cur节点为根，所有子节点(无向树中，剩余节点皆是cur的子节点)到cur的距离和，sz[cur]表示cur所在子树的节点数目，所以,neb[cur]表示cur的子孙节点集合。那么简单的想，自底向上推一遍就可以得到单个节点的答案，所有节点i都做一遍，就可解。但是试了一下，这样做会超时，于是就照题解的思路，再做一遍dfs,以第一遍dfs的解集为基础，每访问一个当前cur的子孙节点，就将该子孙节点视为根，这样只需做如下处理:
，即旧根去掉新根的贡献，新根再补上旧根对它的贡献，更新答案之后将根再换回去。这样就用O(1)复杂度完成了一个节点答案的求解，整个问题总复杂度即O(n).
public:
vector<int> res,dp,sz;
vector<vector<int>> tree;
void dfs(int cur, int pre){
sz[cur]=1;
dp[cur]=0;
for(auto& v:tree[cur]){
if(v==pre)continue;
dfs(v,cur);
dp[cur]+=dp[v]+sz[v];
sz[cur]+=sz[v];
}
}
void dfs2(int cur, int pre){
res[cur]=dp[cur];
for(auto& v:tree[cur]){
if(v==pre)continue;
int pu = dp[cur], pv = dp[v];
int su = sz[cur], sv = sz[v];
dp[cur] -= dp[v] + sz[v];
sz[cur] -= sz[v];
dp[v] += dp[cur] + sz[cur];
sz[v] += sz[cur];
dfs2(v, cur);
dp[cur] = pu, dp[v] = pv;
sz[cur] = su, sz[v] = sv;
}
}
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
res.resize(N, 0);
sz.resize(N, 0);
dp.resize(N, 0);
tree.resize(N, {});
for(auto& edge:edges){
tree[edge[0]].emplace_back(edge[1]);
tree[edge[1]].emplace_back(edge[0]);
}
dfs(0,-1);
dfs2(0,-1);
return res;
}
};
题解
dfs
将一棵树拆成两个子树，子树的根节点分别为A、B，显然A、B节点是相邻的。我们可得所有节点到A的距离和ans[A] = A子树中所有节点到A节点的距离和sum[A]+B子树中所有节点到B节点的距离和sum[B]+Bz子树中节点个数cnt[B]，我们简写为ans[A] = sum[A] + sum[B] + cnt[B]；同理，我们可以得到ans[B] = sum[B] + sum[A] + cnt[A]。 两式作差可以得到ans[A] = cnt[B] - cnt[A] + ans[B]，因为cnt[A]+cnt[B] = N的，因此ans[A] = cnt[B] - cnt[A] + ans[B] + cnt[A] - cnt[A] = ans[B] - cnt[A] + (N-cnt[A])。我们不妨把A看作子节点，B看作是父节点。 我们两次dfs，一次获取每个节点作为根节点的子树中节点的个数，即cnt；一次用来获取答案，也就是ans。 详细过程见代码代码
vector<int> visit;
vector<int> cnt; int dfs(int now,int height,unordered_map<int,vector<int>>& tree){ int ans = height; visit[now] = 1; for(int i=0; i<tree[now].size(); i++){ if(visit[tree[now][i]] == 0){ ans += dfs(tree[now][i],height+1,tree); cnt[now] += cnt[tree[now][i]]; } } cnt[now]++; return ans; } //ans[x] = x@ + y@ + cnty //ans[y] = y@ + x@ + cntx //ans[x] = ans[y] + cnty -cntx void dfs2(int now,int parent,int N,vector<int>&ans,unordered_map<int,vector<int>>& tree){ visit[now] = 1; if(parent!=-1){ ans[now] = ans[parent] - cnt[now] + (N-cnt[now]); } for(int i=0; i<tree[now].size(); i++){ if(visit[tree[now][i]] == 0){ dfs2(tree[now][i],now,N,ans,tree); } } } vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { unordered_map<int,vector<int>> tree; vector<int> ans(N,0); visit = vector<int>(N,0); cnt = vector<int>(N,0); for(int i=0; i<edges.size(); i++){ tree[edges[i][0]].push_back(edges[i][1]); tree[edges[i][1]].push_back(edges[i][0]); } ans[0] = dfs(0,0,tree); //统一将0作为树的根节点 visit = vector<int>(N,0); dfs2(0,-1,N,ans,tree); return ans; }
一个简单的树形dp。两边dfs，第一遍统计出所有子节点到当前节点的距离。然后第二遍dfs用父节点更新子节点
class Solution {
public: const static int MAXN = 10005 ; int dp[MAXN][2], son[MAXN]; int from[MAXN*2], to[MAXN*2], next[MAXN*2], head[MAXN] , tot; vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { init() ; for( int i = 0; i < edges.size(); i++ ){ int u = edges[i][0] ; int v = edges[i][1] ; add( u, v ); add( v, u ); } dfs1( 0, -1 ); dfs2( 0, -1 , N ); vector<int> ans ; for( int i = 0 ; i < N; i++ ){ ans.push_back( dp[i][1] ); } return ans ; } void init(){ memset(head, -1, sizeof head ); tot = 0 ; }void add( int u, int v ){
from[tot] = u ; to[tot] = v ; next[tot] = head[u] ; head[u] = tot++ ; }void dfs1( int u, int f ){
dp[u][0] = 0 ; son[u] = 1 ; for( int i = head[u] ; i != -1; i = next[i] ){ int v = to[i] ; if( v == f ) continue ; dfs1( v, u ); dp[u][0] += dp[v][0] + son[v] ; son[u] += son[v] ; } }void dfs2( int u, int f, int n ){
if( f == -1 ){ dp[u][1] = dp[u][0]; } for( int i = head[u]; i != -1; i = next[i] ){ int v = to[i]; if( v == f ) continue ; dp[v][1] = dp[u][1] + ( n - 2 * son[v] ) ; dfs2( v, u, n ) ; } } };
转载地址：http://iwuii.baihongyu.com/