2020_10_6 每日一题树中距离之和

106 阅读 0 评论 70 点赞

我是靠谱客的博主传统黑猫，最近开发中收集的这篇文章主要介绍2020_10_6 每日一题树中距离之和，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

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

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

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

示例 1:

输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]] 输出:
[8,12,6,10,10,10]
解释: 如下为给定的树的示意图：

我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 也就是
1 + 1 + 2 + 2 + 2 = 8。因此，answer[0] = 8，以此类推。
说明: 1 <= N <= 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

一开始想歪到shortest path上面去了…照着算法导论抄了个算法实现了一下：

class Solution {
    int inf = Integer.MAX_VALUE / 2;
    int[][] multiply(int[][] L, int[][] W, int N) {
        int[][] ret = new int[N][];
        for(int i = 0; i < N; ++i) {
            ret[i] = new int[i + 1];//col <= row
        }
        for(int row = 0; row < N; ++row)
        for(int col = 0; col < row; ++col) {
            int cur = inf;
            int k = 0;
            while(k <= col) {
                cur = Math.min(cur, L[row][k] + W[col][k]);
                ++k;
            }
            while(k <= row) {
                cur = Math.min(cur, L[row][k] + W[k][col]);
                ++k;
            }
            while(k < N) {
                cur = Math.min(cur, L[k][row] + W[k][col]);
                ++k;
            }
            ret[row][col] = cur;
        }
        //System.out.println(Arrays.deepToString(ret));
        return ret;
    }

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        //矩阵乘法 不优化
        
        int[][] L = new int[N][];//W[i][j]为i到j的边
        //阶梯矩阵看会不会爆内存
        
        for(int i = 0; i < N; ++i) {
            L[i] = new int[i + 1];//col <= row
            Arrays.fill(L[i], inf);
            L[i][i] = 0;
        }
        for(int[] e : edges) {
            if(e[0] > e[1]) {
                L[e[0]][e[1]] = 1;
            } else {
                L[e[1]][e[0]] = 1;
            }
        }
        for(int i = 1; i < N - 1; i *= 2) {
            L = multiply(L, L, N);
        }
        
        int[] ret = new int[N];
        for(int i = 0; i < N; ++i)
        for(int j = 0; j < N; ++j) {
            int big = Math.max(i, j), small = Math.min(i, j);
            ret[i] += L[big][small];
        }
        return ret;
    }
}

又超时间又超空间…后来想到这是一个树啊，又没有圈dfs遍历不就好了。求解一个到所有孩子的距离非常简单，问题是要求解所有的。想到，在求解一个的时候可以存一些信息，利用这些信息构建其他点位的距离和。在第一遍遍历的时候，通过dfs，存储每个节点的子树上的节点和(还有孙子之类的, 包括自己)，以及到所有孩子的距离和。数据结构为：

class node {
    int index;
    int parent;
    int sum;
    int numOfChilds;//算自己的孩子的个数
    List<Integer> adjs;
    public node(int i) {
        adjs = new ArrayList<>();
        index = i;
    }
}

设距离和为sum，子树点数为num，父子分别为p和c，递推公式为
p.sum = sum(c.sum + c.num)
p.num = sum(c.num) + 1
某个子树贡献的距离和为子树根节点的的距离和加上子树上的节点的个数。个数的递推就比较直接了。
dfs的代码：

void dfs1(node[] nodes, int cur, int parent) {
    node n = nodes[cur];
    n.parent = parent;
    n.numOfChilds = 1;
    for(int child : n.adjs) {
        if(child == parent) {
            continue;
        }
        dfs1(nodes, child, cur);
        n.sum += nodes[child].sum + nodes[child].numOfChilds;
        n.numOfChilds += nodes[child].numOfChilds;
    }
    return;
}

第一遍dfs结束之后，可以直接给出根节点的距离和为root.sum
然后再作一遍遍历，可以是dfs也可以是bfs，总之只要保证先遍历－节点再遍历子节点就好。这边遍历的时候就可以把sum修改为所有节点的距离和了。设当前节点为n，父节点为p，递推公式为：
n.sum = p.sum + N - n.num * 2;
其实应该写成：
p.sum + (N - n.num) - n.num
即子节点的距离和为父节点的距离和减去自己子树的节点数，加上除了自己子树的其他节点数之和。如果把n当成父节点，N - n.num是p所在的子树节点和，转移根节点的时候到这些节点的距离要+1，而袁兰根节点到自己子树上的节点的距离的代价则不用付了，更新的代码：

ret[n.index] = n.sum = p.sum + N - n.numOfChilds * 2;

所有代码如下，第二遍遍历用的bfs不过dfs也可以：

class Solution {
    
    class node {
        int index;
        int parent;
        int sum;
        int numOfChilds;//算自己的孩子的个数
        List<Integer> adjs;
        public node(int i) {
            adjs = new ArrayList<>();
            index = i;
        }
    }

    void dfs1(node[] nodes, int cur, int parent) {
        node n = nodes[cur];
        n.parent = parent;
        n.numOfChilds = 1;
        for(int child : n.adjs) {
            if(child == parent) {
                continue;
            }
            dfs1(nodes, child, cur);
            n.sum += nodes[child].sum + nodes[child].numOfChilds;
            n.numOfChilds += nodes[child].numOfChilds;
        }
        return;
    }

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        node[] nodes = new node[N];
        for(int i = 0; i < N; ++i) {
            nodes[i] = new node(i);
        }
        for(int[] edge : edges) {
            int n1 = edge[0], n2 = edge[1];
            nodes[n1].adjs.add(n2);
            nodes[n2].adjs.add(n1);
        }

        //先dfs sum存到所有孩子的孩子的孩子的距离和
        dfs1(nodes, 0, 0);
        int[] ret = new int[N];
        ret[0] = nodes[0].sum;
        LinkedList<Integer> queue = new LinkedList<>();
        for(int i : nodes[0].adjs) {
            queue.add(i);
        }
        while(!queue.isEmpty()) {
            node n = nodes[queue.removeFirst()], p = nodes[n.parent];
            ret[n.index] = n.sum = p.sum + N - n.numOfChilds * 2;
            for(int i : n.adjs) {
                if(i != p.index)
                    queue.add(i);
            }
        }
        return ret;
    }
}