834. 树中距离之和
给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
思路
如图，节点2到其他所有节点距离之和可以分解为两部分：
节点2内部子树节点到节点2的距离之和。节点2外部节点到节点2的距离之和。
定义dist[i]表示第i个节点内部子树节点到节点i的距离之和，那么dist[2]怎么求呢？
dist[2] += dist[3] + 2，因为节点3子树有两个节点，需要加上2次红色路径。dist[2] += dist[4] + 1，节点4只有一个节点。
因此，还需要定义nodeNum[i]表示节点i所在子树的节点数量，初始化均为1。思路基本形成了，在递归树的同时计算节点数量，同时计算dist，那么递归的base case是什么呢？很容易想到，叶子节点就是base case，此时dist[i] = 0，nodeNum[i] = 1，其只有一条路径连向父节点。只要从叶子节点开始处理，更新dist和nodeNum即可。先处理叶子节点，再处理根节点的递归树方式是什么呢？很明显就是后序遍历，因此代码如下：
后序遍历处理内部子树
void postOrder(int u
, int v
) {
for(auto& neighbor
: graph
[u
]) {
if(neighbor
== v
) {
continue;
}
postOrder(neighbor
, u
);
dist
[u
] += dist
[neighbor
] + nodeNum
[neighbor
];
nodeNum
[u
] += nodeNum
[neighbor
];
}
}
处理好节点2内部子树节点到节点2的距离之和之后，也就是计算好dist之后，还需要计算节点2外部节点到节点2的距离之和。 如图，节点0的dist[0]是正确的，因为其是根节点，可不可以从这个正确的dist[0]推出其他节点正确的dist[i]呢？如图，dist[2]是存储了节点2内部子树节点到节点2的距离，dist[0]是存储了节点0内部子树节点到节点2的距离。
节点2外部子树的全部节点到节点0的距离 + 节点2外部子树全部节点数 = 节点2外部子树的全部节点到节点2的距离节点2内部子树的全部节点带节点0的距离 - 节点2内部字数全部节点数 = 节点2内部子树的全部节点到节点2的距离两者相加，其实用公式来表示就是dist[0] + N - nodeNum[2] - nodeNum[2] = dist[2]
因而可以从根节点不断更新子节点的dist值，从根节点更新子节点，使用的是先序遍历。
先序遍历处理外部子树
void preOrder(int u
, int v
) {
for(auto& neighbor
: graph
[u
]) {
if(neighbor
== v
) {
continue;
}
dist
[neighbor
] = dist
[u
] - nodeNum
[neighbor
] + nodeNum
.size() - nodeNum
[neighbor
];
preOrder(neighbor
, u
);
}
}
综上，只需要后序遍历+先序遍历就可以解决问题。
全部代码
class Solution {
public:
vector
<int> dist
, nodeNum
;
vector
<vector
<int>> graph
;
void postOrder(int u
, int v
) {
for(auto& neighbor
: graph
[u
]) {
if(neighbor
== v
) {
continue;
}
postOrder(neighbor
, u
);
dist
[u
] += dist
[neighbor
] + nodeNum
[neighbor
];
nodeNum
[u
] += nodeNum
[neighbor
];
}
}
void preOrder(int u
, int v
) {
for(auto& neighbor
: graph
[u
]) {
if(neighbor
== v
) {
continue;
}
dist
[neighbor
] = dist
[u
] - nodeNum
[neighbor
] + nodeNum
.size() - nodeNum
[neighbor
];
preOrder(neighbor
, u
);
}
}
vector
<int> sumOfDistancesInTree(int N
, vector
<vector
<int>>& edges
) {
dist
.resize(N
, 0);
nodeNum
.resize(N
, 1);
graph
.resize(N
);
for(auto& edge
: edges
) {
int u
= edge
[0], v
= edge
[1];
graph
[u
].push_back(v
);
graph
[v
].push_back(u
);
}
postOrder(0, -1);
preOrder(0, -1);
return dist
;
}
};
