LeetCode 834.树中距离之和 Sum of distance in tree

给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。

第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

示例 1:

输入: 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

题目虽然是hard级的，但是感觉难度其实不高。属于较为基础的树状$DP$
首先题目是计算以各点开始的所有距离之和。如果直接计算两点距离。最快的方案是弗洛伊德算法，这是$O(N^3)$的，直接挂
下面分析一下思路：
假设某个节点$root$，令$dp[root][0]$代表以当前节点root为根节点，他下面的子树到他的距离和。$dp[root][1]$代表以当前节点$root$为根节点的话，他下面的所有节点个数（包含他自己），对于叶子节点。$dp[root][0]=0,dp[root][1]=1$
另外为了方便处理，推荐固定以0号节点为全局根节点，因为题目保证了树的连通性，所以以哪个点做根节点其实无所谓。而且题目保证了是树，所以不会出现某个节点有多个父节点的可能
不难看出
$dp[root][0]={\sum }_{leaf\in children}(dp[leaf][0]+dp[leaf][1])$
$dp[root][1]=({\sum }_{leaf\in children}dp[leaf][1])+1$
看不懂上面两个式子可以在纸上画棵树推导一下
此时$dp[0][0]$的结果就是0号节点的答案，可以记做$ans[0]$。假设$chil{d}_x$是某个和0相连的节点。如果以$chil{d}_x$为根节点计算距离的话，因为他和0号节点距离为1，那么不难看出，$chil{d}_x$自己所覆盖的所有子节点距离$chil{d}_x$的距离都会少1，而其它的节点距离会增加1，那么刚刚已经计算出$dp[chil{d}_x][1]$其实就是自己的所有子节点数量，所以剩余的节点数量为$N-dp[chil{d}_x][1]$
可以得到方程：$ans[chil{d}_x]=ans[0]-dp[chil{d}_x][1]+(N-dp[chil{d}_x][1])$
化简一下：$ans[chil{d}_x]=ans[0]-2\ast dp[chil{d}_x][1]+N$
这里只是关于0号节点的转移方程，因为最初只计算了$ans[0]$，有了$ans[0]$根据上面的式子可以计算与0相连的所有儿子节点的$ans[chil{d}_x]$计算出所有$ans[chil{d}_x]$后，可以进一步计算孙子节点从而计算整棵树
一般的动态转移方程：
$$ans[lea{f}_i]=ans[root]-2\ast dp[lea{f}_i][1]+N$$
解法的话，可以先用个$dfs$计算出$dp[i][j]$,然后再次使用$dfs$,计算所有$ans$，两次$dfs$都是$O(N)$的

我的代码，80ms:

class Solution{
private:
	vector<int> graph[10005];
	int dp[10005][2];
	int ans[10005];
public:
	void init_graph(int N,vector<vector<int>>& edges)
	{
		for(int i=0;i<N;++i)
		{
			graph[i].clear();
			dp[i][0]=0;dp[i][1]=1;
		}
		for(int i=0;i<edges.size();++i)
		{
			graph[edges[i][0]].push_back(edges[i][1]);
			graph[edges[i][1]].push_back(edges[i][0]);
		}
	}
	pair<int,int> dfs_fill_dp(int root,int father)
	{
		int weight_sum=0,node_sum=1;
		for(int i=0;i<graph[root].size();++i)
		{
			if(graph[root][i]==father)
				continue;
			pair<int,int> result=dfs_fill_dp(graph[root][i],root);
			weight_sum=weight_sum+result.first+result.second;
			node_sum=node_sum+result.second;
		}
		dp[root][0]=weight_sum,dp[root][1]=node_sum;
		return make_pair(dp[root][0],dp[root][1]);
	}
	void fill_ans_array(int root,int father,int N)
	{
		for(int i=0;i<graph[root].size();++i)
		{
			if(graph[root][i]==father)
				continue;
			ans[graph[root][i]]=ans[root]+N-2*dp[graph[root][i]][1];
			fill_ans_array(graph[root][i],root,N);
		}
	}
    vector<int> sumOfDistancesInTree(int N,vector<vector<int>>& edges)
	{
		vector<int> result;
		init_graph(N,edges);
		pair<int,int> root_result=dfs_fill_dp(0,-1);
		ans[0]=root_result.first;
		fill_ans_array(0,-1,N);
		for(int i=0;i<N;++i)
			result.push_back(ans[i]);
		return result;
    }
};