【leetcode834】树中距离之和

class Solution {
public:
    int distsum[150000]={0};
    int sz[100010]={0};
    vector<int> e[100010];
    vector<int> ans;
    void dfs(int node,int father)
    {
        sz[node]=1;
        for(auto x:e[node])
        {
            if(x!=father)
            {
                dfs(x,node);
                sz[node]+=sz[x];
                distsum[node]+=distsum[x]+sz[x];
            }
        }
    }
    void dfs_change_root(int node,int father)
    {
        ans[node]=distsum[node];
        for(auto x:e[node])
        {
            if(x!=father)
            {
                int a=sz[x],b=sz[node],c=distsum[x],d=distsum[node];
                sz[node]-=sz[x];
                distsum[node]-=distsum[x]+sz[x];
                sz[x]+=sz[node];
                distsum[x]+=distsum[node]+sz[node];
                dfs_change_root(x,node);
                sz[x]=a;
                sz[node]=b;
                distsum[x]=c;
                distsum[node]=d;
            }
        }
    }
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        int num_of_edge=edges.size(),root=0;
        for(int i=0;i<num_of_edge;i++)
        {
            e[edges[i][0]].push_back(edges[i][1]);
            e[edges[i][1]].push_back(edges[i][0]);
        }
        //dfs得出以节点0为根的树的所有子节点距根节点的距离之和,用dfs(i)表示以i为根的子树的子节点与i的距离之和
        dfs(root,-1);
        ans.resize(n);
        //换根
        dfs_change_root(root,-1);
        return ans;
    }
};
上一篇:【JS】箭头函数与普通函数(function)的区别是什么?


下一篇:Python继承/重载/父类/子类