[LeetCode]Sum of Distances in Tree

一个无环的无向图，我们以任何一个节点作根都可以得到一个树。假设我们以编号为0的节点为根，得到的树为T0。我们用dp[i]表示其余所有节点到节点i的距离和，那么对于0的其中一个child c，以c为基准，所有节点和c的距离可以分为两个部分：

• 以c为根的子树，其中的所有节点和c的距离比和0的距离少1
• 除此之外，所有其他的节点和c的距离比和0的距离多1
• 如果以c为根的子树size为s，那么c和其他节点的距离和dp[c] = dp[0] - s + N - s，N为总节点数
• 求每个子树size的话，也是树形dp的思路，对于当前节点root和其其节点c1, c2...，size[root] = 1 + sum(size[ci])

我们dfs两次，第一次求出所有子树的size和其他节点到0的距离。第二次求每个节点到其余节点的距离和。另外可以注意到的一点，如果带权的话方法是一样的，我们知道当前节点到每个子节点的边的长度，对递推方程稍加修改即可，例题可以参考Fermat Point in Graph。对时间复杂度O(N)，空间复杂度O(N):

	class Solution {
	public:
	vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
	vector<unordered_set<int>> g(N);
	for (const auto& e : edges)
	{
	g[e[0]].insert(e[1]);
	g[e[1]].insert(e[0]);
	}
	vector<int> sz(N, 0);
	int totalDist = 0;
	countSize(-1, 0, g, sz, totalDist, 0);
	vector<int> res(N, 0);
	getDist(N, -1, 0, totalDist, g, sz, res);
	return res;
	}
	private:
	void countSize(int from, int to, vector<unordered_set<int>>& g, vector<int>& sz, int& totalDist, int currDist)
	{
	totalDist += currDist;
	int currSize = 1;
	for (const auto& adj : g[to])
	{
	if (adj != from)
	{
	countSize(to, adj, g, sz, totalDist, currDist + 1);
	currSize += sz[adj];
	}
	}
	sz[to] = currSize;
	}
	void getDist(int N, int from, int to, int currDist, vector<unordered_set<int>>& g, vector<int>& sz, vector<int>& res)
	{
	res[to] = currDist;
	for (const auto& adj : g[to])
	if (adj != from)
	getDist(N, to, adj, currDist - sz[adj] + N - sz[adj], g, sz, res);
	}
	};

view raw Sum of Distances in Tree.cpp hosted with ❤ by GitHub