LeeCode 834 DFS

题意

传送门 LeeCode 834. 树中距离之和

题解

枚举根节点 $DFS$ 复杂度 $O(n^2)$，实际上一次 $DFS$ 已经获取了以各节点为根的子树的信息，利用这些信息进行优化。

设 $sum[i],num[i]$ 分别为节点 $i$ 为根的子树的树中距离和与树中节点个数，$ch$ 为节点 $i$ 的子节点，则有
$$\{\begin{array}{l}sum[i]=\sum (sum[ch]+num[ch])\\ num[i]=\sum num[ch]\end{array}$$ 观察根节点的子节点 $ch$，只需要其父节点与子节点的信息，就可以求出 $sum[ch]$。这是因为根节点记录了整棵树的节点信息，那么只需要减去子节点 $ch$ 的部分，就可以获得整棵树中除了以 $ch$ 为根节点的子树以外的部分的信息。那么更新节点 $ch$ 后，节点 $ch$ 记录了以 $ch$ 为根节点的整棵树的信息。那么按照递归深度 $O(1)$ 不断更新节点为新的根节点即可。设 $p$ 为当前节点的父节点，则有
{ s u m [ c h ] = ∑ { s u m [ c h ] + ( s u m [ p ] − s u m [ c h ] − n u m [ c h ] ) + n u m [ p ] − n u m [ c h ] } n u m [ c h ] = n u m [ p ] \begin{cases}sum[ch]=\sum\{sum[ch]+(sum[p]-sum[ch]-num[ch])+num[p]-num[ch]\}\\ num[ch]=num[p] \\ \end{cases} {sum[ch]=∑{sum[ch]+(sum[p]−sum[ch]−num[ch])+num[p]−num[ch]}num[ch]=num[p]​

那么一次 $DFS$ 预处理出以各节点为根节点的子树路径和与节点数；再一次 $DFS$ 更新根节点，求解答案即可。

class Solution
{
public:
    void pre_dfs(int v, int p, vector<vector<int>> &G, vector<int> &sum, vector<int> &num)
    {
        int s = 0, n = 1;
        for (int i = 0; i < G[v].size(); i++)
        {
            int u = G[v][i];
            if (u != p)
            {
                pre_dfs(u, v, G, sum, num);
                s += sum[u] + num[u], n += num[u];
            }
        }
        sum[v] = s, num[v] = n;
    }
    void dfs(int v, int p, vector<vector<int>> &G, vector<int> &sum, vector<int> &num)
    {
        if (p != -1)
        {
            sum[v] = sum[p] + num[p] - (num[v] << 1);
            num[v] = num[p];
        }
        for (int i = 0; i < G[v].size(); i++)
        {
            int u = G[v][i];
            if (u != p)
            {
                dfs(u, v, G, sum, num);
            }
        }
    }
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>> &edges)
    {
        vector<vector<int>> G(N);
        for (auto e : edges)
        {
            G[e[0]].push_back(e[1]);
            G[e[1]].push_back(e[0]);
        }
        vector<int> sum(N), num(N);
        int root = 0;
        pre_dfs(root, -1, G, sum, num);
        dfs(root, -1, G, sum, num);
        return sum;
    }
};