834. Sum of Distances in Tree(leetcode, hard)

https://leetcode.com/problems/sum-of-distances-in-tree/

例如:无向图如下 (参考: https://leetcode.com/problems/sum-of-distances-in-tree/discuss/130567/Two-traversals-O(N)-python-solution-with-Explanation)
在这里插入图片描述
灰色 gray
绿色 green

count(green) 代表 B下面的所有节点的数量 (也可以包括B)
count(gray) 代表A节点的的连接 除了B 之外的所有连接的 节点数量 (也可以包括A)

那么有:

count(green) + count(gray) = N (N表示节点数目)

sum(A) is the sum of the distances between node A and all other nodes. (A到其它所有节点的距离之和)

假设:A 到 B(包括B下面的所有节点)的距离只和 为 sumA(B), 剩下的为sumA(other)

sum(B) is the sum of the distances between node A and all other nodes.(B到其它所有节点的距离之和)

假设:B 到 B下面的所有节点的距离只和 为 sumB(Bchildren), 剩下的即到A(包括A下面的所有节点的)距离只和为sumB(A)

显然有如下的一些等式成立

sumA(B) = sumB(Bchildren) + count(green) + 1 (其中 1 表示 B)
sumB(A) = sumA(other) + count(gray) + 1  (其中1 表示 A)
且
sum(B) = sumB(Bchildren) + sumB(A)

则
sum(A) 
= sumA(B) + sumA(other)
= sumB(Bchildren) + count(green) +  1 + sumB(A) -  count(gray) - 1 
= sumB(Bchildren) + sumB(A) + count(green) - count(gray) 
= sum(B) + count(green) - count(gray) 

或
sum(B) 
= sum(A) + count(gray) - count(green)
// 当 count(green) 为 AB断开连接,B构成的子树的节点数目 也有如下的等式成立
= sum(A) + N - 2 * count(green)

count(gray), count(green) 也是依附于 AB 之间关系而存在的

count(gray) 可以对A叫做 count_direct_link_except_B (与A直接连接但排除B的节点数量)
count(green) 可以对B叫做count_direct_link_except_A(与B直接连接但但排除A的节点数量)

可以看到,其实两个count的含义是一致的


  • 节点关系构成的树
0为root:

    0
   / \
  1   2
     / | \
    3  4  5

1为root:

  1
   \
    0
     \
      2
     / | \
    3  4  5
  • count 表格 (count(i,j)表示去掉i-j的联系后,i的子树节点数目)

如果不包括自己则表格如下:

- 0 1 2 3 4 5
0 - 4 1 - - -
1 0 - - - - -
2 3 - - 4 4 4
3 - - 0 - - -
4 - - 0 - - -
5 - - 0 - - -

如果包括自己则表格如下:

- 0 1 2 3 4 5
0 - 5 2 - - -
1 1 - - - - -
2 4 - - 5 5 5
3 - - 1 - - -
4 - - 1 - - -
5 - - 1 - - -
  • ans

8 12 6 10 10 10

例如:

count(0, 1) = 5
count(1, 0) = 1

sum(1) = sum(0) + count(0, 1) - count(1,0)

sum(1) = sum(0) + 4

// 如果 count(1,0) 表示 0-1连接断开后,1的子树的节点数量(包括1再内)
sum(1) = sum(0) + N - 2 * count(1,0)

sum(1) = sum(0) + 6 - 2 * 1 = sum(1) = sum(0) + 4


2次dfs ac代码 (Java版)

class DataNode{
    int sum; // sum 表示 以这个节点为根的子树的所有子节点到该节点的距离的和, 如果只有一个节点自己,那么,sum=0
    int num; // num 表示 以这个节点为根的子树的所有节点个数,包括该节点,即只有一个节点自己,num = 1
    List<Integer> children;
    boolean isVis;

    public DataNode() {
    }

    public DataNode(int sum, int num, List<Integer> children, boolean isVis) {
        this.sum = sum;
        this.num = num;
        this.children = children;
        this.isVis = isVis;
    }
}

public class Solution {

    final static int MAXN = 10005 ;
    List<DataNode> graph = new ArrayList<>(MAXN);
    List<Integer> ans = new ArrayList<>(MAXN);

    /**
     * 与父节点断开连接后的num
     * @param val
     */
    void dfs(int val){
        graph.get(val).isVis = true;
        for(int valChild: graph.get(val).children){
            if(!graph.get(valChild).isVis){
                dfs(valChild);
                graph.get(val).num += graph.get(valChild).num;
                graph.get(val).sum += graph.get(valChild).sum + graph.get(valChild).num;
            }
        }
    }

    /**
     * 用父节点去更新子节点
     * @param val
     * @param N
     */
    void dfs2(int val, int N){
        for(int valChild: graph.get(val).children){
            // 以此判断是否求解出来
            if(ans.get(valChild) == 0){
                ans.set(valChild, ans.get(val) + N - 2 * graph.get(valChild).num);
                dfs2(valChild, N);
            }
        }
    }

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        for(int i=0;i<N;i++){
            graph.add(new DataNode(0, 1, new ArrayList<>(), false));
            ans.add(0);
        }

        int cnt = edges.length;
        for (int i = 0; i < cnt; i++) {
            int A = edges[i][0];
            int B = edges[i][1];
            graph.get(A).children.add(B);
            graph.get(B).children.add(A);
        }

        int root = 0;
        dfs(root);

        ans.set(root, graph.get(root).sum);
        dfs2(root, N);

        int[] ans2 = new int[N];
        for(int i=0;i<N;i++){
            ans2[i] = ans.get(i);
        }
        return ans2;
    }
}

或者

class DataNode{
    int sum; // sum 表示 以这个节点为根的子树的所有子节点到该节点的距离的和, 如果只有一个节点自己,那么,sum=0
    int num; // num 表示 以这个节点为根的子树的所有节点个数,包括该节点,即只有一个节点自己,num = 1
    List<Integer> children;
    boolean isVis;

    public DataNode() {
    }

    public DataNode(int sum, int num, List<Integer> children, boolean isVis) {
        this.sum = sum;
        this.num = num;
        this.children = children;
        this.isVis = isVis;
    }
}

public class Solution {

    final static int MAXN = 10005 ;
    List<DataNode> graph = new ArrayList<>(MAXN);

    /**
     * 与父节点断开连接后的num
     * @param val
     */
    void dfs(int val){
        graph.get(val).isVis = true;
        for(int valChild: graph.get(val).children){
            if(!graph.get(valChild).isVis){
                dfs(valChild);
                graph.get(val).num += graph.get(valChild).num;
                graph.get(val).sum += graph.get(valChild).sum + graph.get(valChild).num;
            }
        }
    }

    /**
     * 用父节点去更新子节点
     * @param val
     * @param N
     */
    void dfs2(int val, int N){
        for(int valChild: graph.get(val).children){
            // 以此判断是否求解出来
            if(!graph.get(valChild).isVis){
                graph.get(valChild).sum =  graph.get(val).sum + N - 2 * graph.get(valChild).num;
                graph.get(valChild).isVis = true;
                dfs2(valChild, N);
            }
        }
    }

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        for(int i=0;i<N;i++){
            graph.add(new DataNode(0, 1, new ArrayList<>(), false));
        }

        int cnt = edges.length;
        for (int i = 0; i < cnt; i++) {
            int A = edges[i][0];
            int B = edges[i][1];
            graph.get(A).children.add(B);
            graph.get(B).children.add(A);
        }

        int root = 0;
        dfs(root);

        for(int i=0;i<N;i++){
            graph.get(i).isVis = false;
        }
        graph.get(root).isVis = true;
        dfs2(root, N);

        int[] ans2 = new int[N];
        for(int i=0;i<N;i++){
            ans2[i] = graph.get(i).sum;
        }
        return ans2;
    }
}
发布了441 篇原创文章 · 获赞 110 · 访问量 57万+

猜你喜欢

转载自blog.csdn.net/qq_26437925/article/details/93356316