树中距离之和

class Solution {

    List<Integer>[] graph;

    public int[] sumOfDistancesInTree(int n, int[][] edges) {

        this.n = n;

        graph = new List[n];

        for (int i = 0; i < n; i++) {

            graph[i] = new ArrayList<>();

        }

        for (int[] edge : edges) {

            graph[edge[0]].add(edge[1]);

            graph[edge[1]].add(edge[0]);

        }

        sum = new int[n];

        cnt = new int[n];

        ans = new int[n];

        dfs(0, -1);

        ans[0] = sum[0];

        reRoot(0, -1);

        return ans;

    }

    int n;

    //0 为根节点，子树 i 到其所有子节点的距离之和

    int[] sum;

    //0 为根节点，子树 i 及其孩子节点的个数

    int[] cnt;

    int[] ans;

    public int[] dfs(int x, int fa) {

        List<Integer> list = graph[x];

        int curSum = 0;

        int curCnt = 0;

        for (Integer y : list) {

            if (y != fa) {

                int[] dfs = dfs(y, x);

                curCnt += dfs[1];

                curSum += dfs[0];

            }

        }

        sum[x] = curSum + curCnt;

        cnt[x] = curCnt + 1;

        return new int[]{sum[x], cnt[x]};

    }

    public void reRoot(int u, int fa) {

        //u -> v 的转换 =>  v -> u

        // 与 u 为跟节点的不包含 v 的其他节点全部 + 1

        //u 为根节点，到子树 v 的所有距离全部 -1,

        /* 例如

               a1       b1

                \     /

                 u - v - b3 - b4

                /     \

               a2      b2

            现在是 u 转化为 v

            dis [a1], dis [a2] 全部加 1

            dis [b1], dis [b2], dis [b3], dis [b4] 全部减 1

            然后对其求和

         可以等价转化 :

            以 v 为根节点 sum [u] = 以 u 为根节点 sum [u] - (以 u 为根节点 sum [v] + v 及其子节点个数 (因为 v -> u 也有 1 d))

            以 v 为根节点 sum [v] = 以 u 为根节点 sum [v] + 以 v 为根节点 sum [u] + 除 v 之外的包含 u 的相邻节点个数

            ==========================================

            同时也需要修改不同节点为根节点其子树的数量

            以 v 为根节点 cnt [u] = 以 u 为根节点 cnt [u] - 以 u 为根节点 cnt [v]

            以 v 为根节点 cnt [v] = 以 u 为根节点 cnt [u]

         */

        List<Integer> list = graph[u];

        for (Integer v : list) {

            if (v == fa) {

                continue;

            }

            int tf = sum[u];

            int tu = sum[v];

            int tcf = cnt[u];

            int tcu = cnt[v];

			// 除 v 之外的包含 u 的相邻节点个数    

            int count = n - cnt[v];

            // 以 v 为根节点 sum [u] = 以 u 为根节点 sum [u] - (以 u 为根节点 sum [v] + v 及其子节点个数)

            sum[u] = sum[u] - sum[v] - cnt[v];

            // 以 v 为根节点 sum [v] = 以 u 为根节点 sum [v] + 以 v 为根节点 sum [u] + 除 v 之外的包含 u 的相邻节点个数

            sum[v] = sum[v] + sum[u] + count;

            // 以 v 为根节点 cnt [u] = 以 u 为根节点 cnt [u] - 以 u 为根节点 cnt [v]

            cnt[u] = cnt[u] - cnt[v];

            // 以 v 为根节点 cnt [v] = 以 u 为根节点 cnt [u]

            cnt[v] = tcf;

            ans[v] = sum[v];

            reRoot(v, u);

            sum[u] = tf;

            sum[v] = tu;

            cnt[u] = tcf;

            cnt[v] = tcu;  

        }

    }

}