834. Sum of Distances in Tree c++

最新推荐文章于 2025-01-13 15:16:10 发布

原创最新推荐文章于 2025-01-13 15:16:10 发布 · 225 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#leetcode

leetcode 专栏收录该内容

155 篇文章

订阅专栏

给定一棵无向、连通的树，节点从1到n编号。第i条边将节点u和v连接在一起。返回一个列表，其中第i个元素是节点i与其他所有节点之间的距离之和。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

An undirected, connected tree with N nodes labelled 0...N-1 and N-1edges are given.

The ith edge connects nodes edges[i][0] and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node i and all other nodes.

Example 1:

Input: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: 
Here is a diagram of the given tree:
  0
 / \
1   2
   /|\
  3 4 5
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.  Hence, answer[0] = 8, and so on.

Note: 1 <= N <= 10000

class Solution {
public:
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) 
    {
        if(edges.empty()) return {0};
        v.resize(N);
        ans.resize(N,0);
        for(int i=0;i<edges.size();i++)
        {
            v[edges[i][0]].connect.push_back(edges[i][1]);
            v[edges[i][1]].connect.push_back(edges[i][0]);
        }
        dfs(edges[0][0]);
        ans[edges[0][0]] = v[edges[0][0]].sum;
        dfsans(N,edges[0][0]);
        return ans;
        
    }
    struct TreeNode
    {
        int sum=0;
        int num=1;
        vector<int> connect;
        bool visited=false;
    };
    vector<TreeNode> v;
    vector<int> ans;
    void dfs(int val)
    {
        v[val].visited = true;
        for(int i=0;i<v[val].connect.size();i++)
        {
            if(!v[v[val].connect[i]].visited)
            {
                dfs(v[val].connect[i]);
                v[val].num += v[v[val].connect[i]].num;
                v[val].sum += v[v[val].connect[i]].sum + v[v[val].connect[i]].num;
            }
        }
    }
    void dfsans(int N,int val)
    {
        for(int i=0;i<v[val].connect.size();i++)
        {
            if(ans[v[val].connect[i]] == 0)
            {
                ans[v[val].connect[i]] = ans[val] + N - 2*v[v[val].connect[i]].num;
                dfsans(N,v[val].connect[i]);
            }
        }
    }
};