51nod 1322 关于树的函数

题目
题解

#include<bits/stdc++.h>
#include<tuple>
using namespace std;
typedef tuple<int,int> ti;
typedef long long ll;
#define mt make_tuple
const int N=4002;
struct node{
    int to,ne;
}e1[N<<1],e2[N<<1];
int n,a,b,x,y,xx[N],yy[N],h1[N],h2[N],tot1,tot2,i,vis[N];
ll ans;
inline char gc(){
    static char buf[100000],*p1=buf,*p2=buf;
    return p1==p2&&(p2=(p1=buf)+fread(buf,1,100000,stdin),p1==p2)?EOF:*p1++;
}
#define gc getchar
inline int read(){
    int x=0,fl=1;char ch=gc();
    for (;ch<48||ch>57;ch=gc())if(ch=='-')fl=-1;
    for (;48<=ch&&ch<=57;ch=gc())x=(x<<3)+(x<<1)+(ch^48);
    return x*fl;
}
void add1(int x,int y){
    e1[++tot1]=(node){y,h1[x]};
    h1[x]=tot1;
}
void add2(int x,int y){
    e2[++tot2]=(node){y,h2[x]};
    h2[x]=tot2;
}
void dfs1(int u,int fa){
    vis[u]=1;a++;
    for (int i=h1[u],v;i;i=e1[i].ne)
        if ((v=e1[i].to)!=fa) dfs1(v,u);
}
ti dfs2(int u,int fa){
    int cnt=0,b=1,x,y;
    if (vis[u]) cnt++;
    for (int i=h2[u],v;i;i=e2[i].ne)
        if ((v=e2[i].to)!=fa){
            tie(x,y)=dfs2(v,u);
            cnt+=x;b+=y;
        }
    if (b==n) return mt(cnt,b);
    int mx=max(max(max(cnt,a-cnt),b-cnt),n-a-b+cnt);
    ans+=mx*mx;
    return mt(cnt,b);
}
int main(){
    n=read();
    for (i=1;i<n;i++){
        x=read();y=read();
        add1(x,y);add1(y,x);
        xx[i]=x;yy[i]=y;
    }
    for (i=1;i<n;i++){
        x=read();y=read();
        add2(x,y);add2(y,x);
    }
    for (i=1;i<n;i++){
        a=0;
        memset(vis,0,sizeof(vis));
        dfs1(xx[i],yy[i]);
        dfs2(1,-1);
    }
    printf("%lld",ans);
}