2023-07-16 LeetCode每日一题（树中距离之和）

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 * 10⁴
• 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有x₁，x₂，x₃个子节点，那么dp[u]就应当等于3+dp[x₁]+dp[x₂]+dp[x₃]。那么sz[u]表示sz[x₁] + sz[x₂] + sz[x₃]，这个是dfs1中所完成的工作。当然一开始所有的sz都应该为1。我们按照这样的思路进行n次，那么一定能解决问题，但是时间复杂度很高，能不能优化呢。

(4) 经过一次树形动态规划后其实我们获得了在 u 为根的树中，每个节点为根的子树的答案 dp，我们可以利用这些已有信息来优化时间复杂度。本来我们所求的是以u为根节点，v为子节点，后面我们要进行置换一下，那么v是根节点，u为子节点了。除了dp[u]和dp[v]的值，其他的值并不需要变化。首先u要减去v的贡献。v要加上u的贡献。这样我们便实现了换根的操作了。