Atcoder[ABC401F] Add One Edge 3 题解

[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)有两种情况：

1. 经过边( i , j ) (i,j)(i,j)：a i + b j + 1 a_i+b_j+1ai​+bj​+1。
2. 不经过边( 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(nlog​n)，瓶颈在于排序。
注意到可以使用桶排序，时间复杂度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;}