树上求和(DFS)

树上求和(DFS)

传送门

思路:计算每条边的贡献次数,然后排序即可。
每条边贡献次数公式: c n t = s i z e [ u ] × ( n s i z e [ u ] ) cnt=size[u]\times (n-size[u])

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=1e5+5;
#define mst(a) memset(a,0,sizeof a)
#define lx x<<1
#define rx x<<1|1
#define reg register
#define PII pair<int,int>
#define fi first 
#define se second
inline void read(int &x){ 
	x=0;int w=1;
	char ch=getchar();
	while(ch<'0'||ch>'9') {if(ch=='-') w=-1;ch=getchar();}
	for(;ch>='0'&&ch<='9';ch=getchar())
		x=(x<<3)+(x<<1)+(ch&15);
	x*=w; 
}
int sz[N],n;
vector<int>e[N];
vector<ll>ans;
void dfs(int u,int fa){
	//printf("u=%d,fa=%d\n",u,fa); 
     sz[u]=1;
    for(auto v:e[u]){
        if(v==fa) continue;
        dfs(v,u);
        sz[u]+=sz[v];
    }
    ans.push_back(1LL*sz[u]*(n-sz[u]));//计算以u为终点的每条边的贡献次数. 所有边只遍历一次. 
}
int main(){
    read(n);
    for(int i=1,u,v;i<n;i++)
        read(u),read(v),e[u].push_back(v),e[v].push_back(u);
    dfs(1,0);
    sort(ans.begin(),ans.end());
    ll res=0;
    for(int i=1;i<n;i++){//注意这里从1开始,因为从根为终点的边(没有这个边)贡献是0. 
        res+=(n-i)*ans[i];
    }
    printf("%lld\n",res);
	return 0;
}

猜你喜欢

转载自blog.csdn.net/weixin_45750972/article/details/106659318