[ABC401F] Add One Edge 3
思路
设第一棵树的直径长度为l 1 l1l1，第二棵树的直径长度为l 2 l2l2，a i a_iai为第一棵树中以点i ii为端点的路径的长度最大值，b i b_ibi为第二棵树中以点i ii为端点的路径的长度最大值。则f ( i , j ) f(i,j)f(i,j)有两种情况：
- 经过边( i , j ) (i,j)(i,j)：a i + b j + 1 a_i+b_j+1ai+bj+1。
- 不经过边( i , j ) (i,j)(i,j)：max ( l 1 , l 2 ) \max (l1,l2)max(l1,l2)。
两种情况取较大值即可。
根据直径的性质，a i a_iai与b i b_ibi一定经过直径的其中一个端点，故在求直径的时候可以顺便求出a i a_iai和b i b_ibi。
我们将a aa与b bb排序。对于每个i ii，可以二分/双指针求出b i + a i ≤ max ( l 1 , l 2 ) b_i+a_i\le \max (l1,l2)bi+ai≤max(l1,l2)的b i b_ibi，O ( 1 ) O(1)O(1)快速计算答案。
代码
时间复杂度O ( n log n ) O(n\log_{}{n} )O(nlogn)，瓶颈在于排序。
注意到可以使用桶排序，时间复杂度O ( n ) O(n)O(n)。
#include<bits/stdc++.h>usingnamespacestd;constintN=2e5+5;vector<int>e[N];intn,k,a[N],b[N],dep[N],fa[N];voiddfs(intx){dep[x]=dep[fa[x]]+1;if(dep[x]>dep[k])k=x;for(inti=0;i<e[x].size();i++)if(e[x][i]!=fa[x]){fa[e[x][i]]=x;dfs(e[x][i]);}}boolbj[N];intl[N],cnt;voiddfss(intx,inty)//求a,b{a[x]=y,bj[x]=1;for(inti=0;i<e[x].size();i++)if(e[x][i]!=fa[x]&&!bj[e[x][i]]){fa[e[x][i]]=x;dfss(e[x][i],y+1);}}intsolve(){cin>>n;for(inti=1;i<n;i++){intu,v;cin>>u>>v;e[u].push_back(v),e[v].push_back(u);}memset(bj,0,sizeof(bj));k=cnt=fa[1]=0;dfs(1);intkkk=k;k=fa[kkk]=0;dfs(kkk);while(k)l[++cnt]=k,bj[k]=1,k=fa[k];//求直径for(inti=1;i<=cnt;i++){fa[l[i]]=0;dfss(l[i],max(i-1,cnt-i));}for(inti=1;i<=n;i++)e[i].clear();returncnt-1;}inttong[N];voidsor()//神秘桶排序{memset(tong,0,sizeof(tong));for(inti=1;i<=n;i++)tong[a[i]]++;intcnt=0;for(inti=1;i<=2e5;i++)while(tong[i]--)a[++cnt]=i;}intmain(){ios::sync_with_stdio(0);cin.tie(0),cout.tie(0);intl1=solve(),m=n;sor();for(inti=1;i<=n;i++)b[i]=a[i];intl2=solve();sor();intnow=0;longlongans=0,sum=0;for(inti=1;i<=m;i++)sum+=b[i]+1;for(inti=n;i>=1;i--)//注意枚举顺序{while(now+1<=m&&b[now+1]+1+a[i]<=max(l1,l2))now++,sum-=b[now]+1;ans+=now*1ll*max(l1,l2);ans+=sum+1ll*(m-now)*a[i];}cout<<ans;return0;}
