每日一题：树中距离之和

题目

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

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

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

Solution

首先只考虑查询一个点的情况，容易得到状态转移：$dp[u] = \sum_{v \in son[u]}{dp[v] + sz[v]}$。

$son[u]$ 表示 $u$ 的所有子节点，$sz[v]$ 表示以 $v$ 为根的子树大小，$dp[u]$ 表示以 $u$ 为根的子树，它的所有子节点到它的距离之和。

那么 $dp[u]$ 就是所有子节点的 $dp[v]$ + 子节点对应的子树大小个点从 $u$ 到 $v$ 的距离。

然后考虑换根，根从 $u$ 变为子节点 $v$，以 $u$ 为根的子树大小减少 $sz[v]$，$dp[u]$ 也要减去 $v$ 的贡献：

$sz[u] -= sz[v], dp[u] -= dp[v] + sz[v]$

然后 $v$ 做为根，获得了 $u$ 结点的贡献，同时以 $v$ 为根的子树大小增大 $sz[u]$：

$sz[v] += sz[u], dp[v] += dp[u] + sz[u]$

同时注意：递归在进行回溯的时候，需要恢复现场，否则在计算兄弟结点时，维护的 $dp$ 和 $sz$ 将是错误的。

需要两次dfs，时间复杂度为 $O(n)$。

Code

class Solution {
public:
    vector<int> g[10005];
    vector<int> res;
    int dp[10005], sz[10005];

    void dfs1(int u, int fa) {
        sz[u] = 1;
        dp[u] = 0;
        for (auto v: g[u]) {
            if (v == fa) continue;
            dfs1(v, u);
            sz[u] += sz[v];
            dp[u] += dp[v] + sz[v];
        }
    }

    void dfs2(int u, int fa) {
        res[u] = dp[u];
        for (auto v: g[u]) {
            if (v == fa) continue;
            int du = dp[u], dv = dp[v];
            int su = sz[u], sv = dp[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] = du, dp[v] = dv;
            sz[u] = su, sz[v] = sv;
        }
    }

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        for (auto edge: edges) {
            g[edge[0]].push_back(edge[1]);
            g[edge[1]].push_back(edge[0]);
        }
        res.resize(N, 0);
        dfs1(0, -1);
        dfs2(0, -1);
        return res;
    }
};