2023-07-16每日一题
一、题目编号
834. 树中距离之和
二、题目链接
点击跳转到题目位置
三、题目描述
给定一个无向、连通的树。树中有 n 个标记为 0…n-1 的节点以及 n-1 条边 。
给定整数 n 和数组 edges ， edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。
返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。
提示:
- 1 <= n <= 3 * 104
- edges.length == n - 1
- edges[i].length == 2
- 0 <= ai, bi < n
- ai != bi
- 给定的输入保证为有效的树
四、解题代码
class Solution {int dp[30004];//表示dp[u] 表示以u为根的子树,它的所有子节点到它的距离之和int sz[30004];//表示以u为根的子树的节点数量vector<int> ans;//表示最终的答案vector<int>Edge[30004];void dfs1(int u, int f) {sz[u] = 1;dp[u] = 0;for (int i = 0; i < Edge[u].size(); ++i) {int v = Edge[u][i];if (v == f) {continue;}dfs1(v, u);dp[u] += dp[v] + sz[v];sz[u] += sz[v];}}void dfs2(int u, int f) {ans[u] = dp[u];for (int i = 0; i < Edge[u].size(); ++i) {int v = Edge[u][i];if (v == f) {continue;}int pu = dp[u], pv = dp[v];int su = sz[u], sv = sz[v];dp[u] -= dp[v] + sz[v];sz[u] -= sz[v];dp[v] += dp[u] + sz[u];sz[v] += sz[u];dfs2(v, u);dp[u] = pu, dp[v] = pv;sz[u] = su, sz[v] = sv;}}public:vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {for(int i = 0; i < n; ++i){Edge[i].clear();}for(int i = 0; i < n - 1; ++i){int x = edges[i][0];int y = edges[i][1]; Edge[x].push_back(y);Edge[y].push_back(x);}memset(dp, 0, sizeof(dp));memset(sz, 0, sizeof(sz));ans.resize(n);dfs1(0, -1);dfs2(0, -1);return ans;}
};
五、解题思路
(1) 采用树型dp来解决该问题。
(2) 先设置初始状态，dp[u] 表示以u为根的子树,它的所有子节点到它的距离之和。sz[u]表示以u为根的子树的节点数量。vector ans表示最终答案。
(3) 那么需要问一个问题，我们所需要求的是dp[u]是怎么得到的。节点u有x1，x2，x3个子节点，那么dp[u]就应当等于3+dp[x1]+dp[x2]+dp[x3]。那么sz[u]表示sz[x1] + sz[x2] + sz[x3]，这个是dfs1中所完成的工作。当然一开始所有的sz都应该为1。我们按照这样的思路进行n次，那么一定能解决问题，但是时间复杂度很高，能不能优化呢。
(4) 经过一次树形动态规划后其实我们获得了在 u 为根的树中，每个节点为根的子树的答案 dp，我们可以利用这些已有信息来优化时间复杂度。本来我们所求的是以u为根节点，v为子节点，后面我们要进行置换一下，那么v是根节点，u为子节点了。除了dp[u]和dp[v]的值，其他的值并不需要变化。首先u要减去v的贡献。v要加上u的贡献。这样我们便实现了换根的操作了。
