LeetCode 834. 树中距离之和 (树形DP、换树根)

树中距离之和
这道题目，很容易想到$O(n^2)$的做法。

• 如何只求出到一个点的距离:
  状态：$f[x]$就表示所有的点到x的距离；
  DP方程: $f[x]=f[y_1]+f[y_2]+\dots \dots +f[y_k]+\dots \dots +s[y_1]+s[y_2]+\dots \dots +s[y_k]+\dots \dots$
  其中，$y_1、y_2、y_k$等时$x$的所有孩子。
  $s[x]$表示以$x$为根的子树的包含的所有节点的数量（边权为1）。
  时间复杂度：$O(n)$

• 优化：如果如下图一样，$x$和一个孩子$y$交换了一下，y成了树根，发现绝大数节点的状态值并没有发生变化。仅仅是它们两个节点的值发生了变化。
  于是，可以同样按照遍历树的顺序，把每一个节点都当成树根。换根的时候只发生在父-子之间。

• 实现的注意点：
  由于是无向图到有根树的转变，在遍历的时候需要记录上一个节点，防止“反复横跳”。
  其次，换根的时候，需要回溯现场，以便x的其他孩子成为树根。

const int N = 10010;
class Solution {
    
    
public:
    vector<int> g[N];
    int f[N] ,s[N];
    vector<int> ans;
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
    
    
        ans.resize(n,-1);
        // memset(f,-1,sizeof(f));
        for(auto&v:edges){
    
    
            g[v[0]].push_back(v[1]);
            g[v[1]].push_back(v[0]);
        }
        dfs(0,-1);
        dp(0,-1);
        return ans;
    }

    void dp(int x,int father){
    
    
        ans[x] = f[x];
        for(int y:g[x]){
    
    
            if(y==father) continue;
            int oldfx = f[x],oldsx = s[x];
            int oldfy = f[y],oldsy = s[y]; 
            f[x] -= f[y]+s[y];
            s[x] -= s[y]; 
            f[y] += f[x]+s[x];
            s[y] += s[x];
            dp(y,x);
             // 还原现场
            f[x] = oldfx;
            s[x] = oldsx;
            f[y] = oldfy;
            s[y] = oldsy;
           
        }
    }

    void dfs(int x,int father){
    
    
        f[x] = 0;
        s[x] = 1;
        for(int y:g[x]){
    
    
            if(y==father) continue;
            dfs(y,x);
            f[x] += f[y]+s[y];
            s[x] += s[y];
        }
    }
};