牛客小白月赛23 G. 树上求和

Link
题意：
树的每条边上都有一个正整数，这个正整数被称作这条边的颜色，规定一条树链的权值为这条树链上所有边的颜色的代数和
而整棵树的权值为所有不同的树链的权值的代数和
已知所有边的颜色集合恰好为 \(1\) 到 \(n-1\) 这 \(n-1\) 个不同的正整数，请你为每条边安排一种颜色，使得这棵树的权值尽量小，求出这个最小的权值
思路：
设一条边上深度较大的那个点为 \(x\),点 \(x\) 的子树大小为 \(sz[x]\)，则一条边所在链个数 \(sz[x]*(n-sz[x])\)
给所在链个数较大的分配较小的值
代码：

#include<bits/stdc++.h>
 
using namespace std;
 
const int N=1e5+5;
int n;
int cnt,to[N*2],nxt[N*2],head[N];
long long sz[N],tot[N];
 
void addedge(int a,int b) {
    cnt++;
    to[cnt]=b;
    nxt[cnt]=head[a];
    head[a]=cnt;
}
void dfs(int u,int pre) {
    sz[u]=1;
    for(int i=head[u];i;i=nxt[i]) {
        int v=to[i];
        if(v==pre) continue;
        dfs(v,u);
        sz[u]+=sz[v];
    }
    tot[u]=sz[u]*(n-sz[u]);
}
int main() {
    //freopen("in.txt","r",stdin);
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin>>n;
    for(int i=1;i<n;i++) {
        int a,b;
        cin>>a>>b;
        addedge(a,b);
        addedge(b,a);
    }
    dfs(1,0);
    sort(tot+2,tot+1+n);
    long long ans=0;
    for(int i=2;i<=n;i++) ans+=tot[i]*(n-i+1);
    cout<<ans<<endl;
    return 0;
}