Comet OJ - Contest #6 C 一道树题 数学 + 推导

Code:

#include <bits/stdc++.h>
#define setIO(s) freopen(s".in","r",stdin)
#define maxn 300000
#define ll long long
using namespace std;
vector<int>G[maxn];
ll sumv[maxn];
ll Sum(ll k) {
    return k*(k+1)/2;
}
ll calc(ll k) {
    return 1ll*(k+1)*Sum(k)-1ll*sumv[k];
}
void Initialize() {
    for(int i=1;i<maxn;++i) sumv[i]=sumv[i-1]+1ll*i*i;
}
int main() {
    // setIO("input");
    int n,m;
    scanf("%d",&n),m=n-1;
    for(int i=1;i<=m;++i) {
        int a,b;
        scanf("%d%d",&a,&b);
        G[a].push_back(i), G[b].push_back(i);
    }
    Initialize();
    ll ans=(ll)((ll)(1ll*m*n+n)*m)/2;
    ans -= (ll) calc(m);
    for(int i=1;i<=n;++i) {
        sort(G[i].begin(), G[i].end());
        for(int j=0;j<G[i].size();++j) {
            if(j==0) {
                ans-=Sum((ll)G[i][j]-1);
            }
            if(j>0) {
                ans-=Sum((ll)G[i][j]-G[i][j-1]-1);
            }
            if(j==G[i].size()-1) {
                ans-=Sum((ll)m-G[i][j]);
            }
        }
    }
    printf("%lld\n",ans);
    return 0;
}