834. 树中距离之和

wion03

已于 2023-07-17 12:05:55 修改

阅读量242

点赞数

分类专栏： LeetCode每日一题系列文章标签：算法数据结构动态规划深度优先遍历 c++

于 2023-07-17 10:58:47 首次发布

本文链接：https://blog.csdn.net/wion03/article/details/131757467

版权

LeetCode每日一题系列专栏收录该内容

11 篇文章 0 订阅

订阅专栏

文章讨论了一种优化算法，用于解决给定无向、连通树中每个节点与其他所有节点距离之和的问题。初始尝试使用深度优先遍历导致超时，随后采用动态规划结合深度优先遍历和记忆化搜索的方法，实现了时间复杂度为O(N)的解决方案，有效地处理了大数据量的情况。

摘要由CSDN通过智能技术生成

题目描述：

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

给定整数 n 和数组 edges ， edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。

返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。

示例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，以此类推。

示例 2:

输入：n = 1，edges = [ ]
输出：[0]

示例3：

输入：n = 2，edges = [ [1,0] ]
输出： [1,1]

提示:

1 <= n <= 3 * $10^{4}$
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
给定的输入保证为有效的树

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

初步尝试：深度优先遍历所有结点

拿到这道题，作者首先想到就是每次从一个节点出发，遍历整个树，计算所有节点到它的距离之和。很显然这种暴力的方法时间复杂度为 $O(N^{2})$ ，而题目中数据范围最大为30000，明显会超时。

这里附上作者初次提交的代码，样例能过但数据量到10000时就超时了，因此我们需要考虑更优化的算法

class Solution {
public:
    struct node
    {
        vector<int> next;
        int nums=0;
    }nodes[30005];
    int tag[30005]={0};  //标记数组，每次遍历前重置，遍历过的结点置为1
    vector<int> ans;
    int out[30005]={0};

    void dfs(int u,int x,int n)  //n是结点x到u的距离，也就是以u为根的树中x的深度
    {
        tag[x]=1;
        for(int i=0;i<=nodes[x].nums-1;i++)
        {
            if(tag[nodes[x].next[i]]==0)
            {
                tag[nodes[x].next[i]]=1;
                dfs(u,nodes[x].next[i],n+1);
            }
        }
        out[u]+=n;  //将当前结点的距离加入答案中
    }

    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        if(n == 1) return {0};
        for(int i=0;i<=n-2;i++)
        {
            int x=edges[i][0];
            int y=edges[i][1];
            nodes[x].nums++;
            nodes[x].next.push_back(y);
            nodes[y].nums++;
            nodes[y].next.push_back(x);
        }
        for(int i=0;i<=n-1;i++)
        {
            memset(tag,0,sizeof(tag));
            dfs(i,i,0);
            ans.push_back(out[i]);
        }
        return ans;
        
    }
};

题解：动态规划+深度优先遍历+记忆化搜索

首先，我们考虑一个节点的情况，即我们只求根节点到所有其它节点的距离之和：

令dp[node]表示以node为根的子树中所有节点到node的距离之和，counts[node]表示以node为根的子树的节点数。（注意此时整个树的结点还是0）
那么我们可以通过观察得到一个公式：

当前结点的dp值等于该结点所有的子节点的dp值和counts值之和： $dp[node]=\sum_{u\in son[node]}^{}dp[u]+counts[u]$

解释一下原理：

以这颗树中的 0->2->(3,4,5) 为例，如果已经知道了 2 这个节点到其所有子节点 3,4,5 的距离之和 dp[2]，那么我们能否通过 2 来得到 0 到 3,4,5 的距离之和呢？

因为 0 到 2 的距离为1，那么 0 到 3,4,5 的距离应该在 2 的基础上加1，而每个子节点都加一就相当于加上这个子树的子节点的数量了，即：

dp[0] = dp[1] +1+dp[2]+1+1+1+1 = dp[1] + counts[1] + dp[2] + counts[2]

通过上述方法，我们从 0 节点开始深度优先遍历，每次给节点赋予初始值：dp =0 ，counts = 1

然后遍历当前结点的所有子节点，最终获得一个以0为根的树的dp和counts状态。

 void dfs1(int father,int current_node)
    {
        dp[current_node]=0;
        counts[current_node]=1; //初次到达结点时的初始化
        for(int i=0;i<=nodes[current_node].nums-1;i++)
        {
            int son_node=nodes[current_node].next[i];
            if(son_node!=father)  //不能向上遍历
            {
                dfs1(current_node,son_node);
                dp[current_node]+=dp[son_node]+counts[son_node];
                counts[current_node]+=counts[son_node];
            }
        }
    }

该方法的时间复杂度为O（N）。

接下来，考虑题目给出的情况：获取所有节点到其他所有节点的距离之和。

如果我们用上面的dfs1方法来遍历所有节点作为根，那么得到的的算法的时间复杂度还是 $O(N^{2})$

显然，我们需要其他的办法。

考虑：如何从当前的以0为根的状态得到以2为根的状态呢？

观察两图，发现只有交换的两个节点 0 和 2 的dp值以及counts发生了变化：

因为 2 成为了 0 的父节点，要从 dp[0] 和 counts[0] 中减去 2 带来的部分：

dp[0] = dp[0] - dp[2] - counts[2]

counts[0] = counts[0] - counts[2]

而 0 成为了 2 的子节点，同样要在 dp[2] 和 counts[2] 中加上 0 带来的部分：

dp[2] = dp[2] + dp[0] + counts[0]

counts[2] = counts[2] + counts[0]

这样，在不改变其他结点状态的情况下，我们用常量级的时间就完成了根的变换操作，获得了新的dp值。

void dfs2(int oldroot,int current_root)
    {
        ans[current_root] = dp[current_root];
        //记录以当前节点为根时的dp值
        for(int i=0;i<=nodes[current_root].nums-1;i++)
        {
            int next_root = nodes[current_root].next[i];
            if(next_root != oldroot) //不要走回头路，年轻人
            {
                //记录此时的状态，因为下次交换时要确保还是以current_node为根时的状态
                int a= dp[current_root];
                int b= dp[next_root];
                int c= counts[current_root];
                int d= counts[next_root];
                //以下是交换操作
                dp[current_root] = dp[current_root] - dp[next_root]-counts[next_root];
                counts[current_root]-=counts[next_root];
                dp[next_root] = dp[next_root]+dp[current_root]+counts[current_root];
                counts[next_root]+=counts[current_root];

                dfs2(current_root,next_root);
                //回溯
                dp[current_root]=a;
                dp[next_root]=b;
                counts[current_root]=c;
                counts[next_root]=d;
            }
        }
    }

这里一定要注意的是：当我们用深度优先搜索遍历，交换了根后进行下次交换前，一定要还原根的状态。

完整代码如下：

class Solution {
public:
    struct node
    {
        vector<int> next;
        int nums=0;
    }nodes[30005];
    vector<int> ans;
    vector<int> dp;
    vector<int> counts;

    void dfs1(int father,int current_node)
    {
        dp[current_node]=0;
        counts[current_node]=1; //初次到达结点时的初始化
        for(int i=0;i<=nodes[current_node].nums-1;i++)
        {
            int son_node=nodes[current_node].next[i];
            if(son_node!=father)  //不能向上遍历
            {
                dfs1(current_node,son_node);
                dp[current_node]+=dp[son_node]+counts[son_node];
                counts[current_node]+=counts[son_node];
            }
        }
    }

    void dfs2(int oldroot,int current_root)
    {
        ans[current_root] = dp[current_root];
        //记录以当前节点为根时的dp值
        for(int i=0;i<=nodes[current_root].nums-1;i++)
        {
            int next_root = nodes[current_root].next[i];
            if(next_root != oldroot) //不要走回头路，年轻人
            {
                //记录此时的状态，因为下次交换时要确保还是以current_node为根时的状态
                int a= dp[current_root];
                int b= dp[next_root];
                int c= counts[current_root];
                int d= counts[next_root];
                //以下是交换操作
                dp[current_root] = dp[current_root] - dp[next_root]-counts[next_root];
                counts[current_root]-=counts[next_root];
                dp[next_root] = dp[next_root]+dp[current_root]+counts[current_root];
                counts[next_root]+=counts[current_root];

                dfs2(current_root,next_root);
                //回溯
                dp[current_root]=a;
                dp[next_root]=b;
                counts[current_root]=c;
                counts[next_root]=d;
            }
        }
    }


    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        dp.resize(n,0);  //dp[node]表示以node为根的树，其所有节点到其距离的和
        counts.resize(n,0);  //counts[son]表示以son为根的子树的结点数
        ans.resize(n,0);  //记录答案

        if(n == 1) return {0};
        for(int i=0;i<=n-2;i++)
        {
            int x=edges[i][0];
            int y=edges[i][1];
            nodes[x].nums++;
            nodes[x].next.push_back(y);
            nodes[y].nums++;
            nodes[y].next.push_back(x);
        }
        
        dfs1(-1,0);  //第一遍遍历以0为根节点获得基础的数据
        dfs2(-1,0);  //第二遍遍历依次以各节点为根获得答案
        
        return ans;
        
    }
};