loj6119 「2017 山东二轮集训 Day7」国王

题目描述

在某个神奇的大陆上，有一个国家，这片大陆的所有城市间的道路网可以看做是一棵树，每个城市要么是工业城市，要么是农业城市，这个国家的人认为一条路径是 exciting 的，当且仅当这条路径上的工业城市和农业城市数目相等。现在国王想把城市分给他的两个儿子，大儿子想知道，他选择一段标号连续的城市作为自己的领地，并把剩下的给弟弟，能够满足两端都是自己城市的 exciting 路径比两端都是弟弟的城市的 exciting 路径数目多的方案数。

输入格式

第一行一个正整数 n n n。
第二行 n n n 个整数依次描述城市的性质，1 1 1 为工业，0 0 0 为农业。
接下来 n−1 n - 1 n−1 行每行两个正整数描述一条道路。

输出格式

输出一个整数表示答案。

样例

样例输入

样例输出

数据范围与提示

n≤100000 n leq 100000 n≤100000

正解：点分治。

对于每个右端点，我们找出极小的左端点使得在这个区间内大儿子不能获利，显然这是满足单调性的。

然后我们设$A$为两个端点都在区间内的路径数量，$B$为两个端点都在区间外的路径数量，如果$A>B$，那么左端点就可以往右移。

设$C$为两个端点分别在区间内外的答案，我们发现$2A+C>2B+C$与前面的不等式是等价的。

设$sum$为总路径数$*2$，$f[i]$为一个端点为$i$的路径数量，那么$2A+C=sum_{i=l}^{r}f[i]$，$2B+C=sum-sum_{i=l}^{r}f[i]$。

然后用点分治来统计一下路径就行了。。

 1 #include <bits/stdc++.h>
 2 #define il inline
 3 #define RG register
 4 #define ll long long
 5 #define N (500005)
 6 
 7 using namespace std;
 8 
 9 struct edge{ int nt,to; }g[N];
10 struct data{ int i,l; }st[N];
11 
12 int head[N],tong[N],vis[N],dis[N],son[N],sz[N],a[N],n,num,top;
13 ll f[N],now,sum,ans;
14 
15 il int gi(){
16   RG int x=0,q=1; RG char ch=getchar();
17   while ((ch<'0' || ch>'9') && ch!='-') ch=getchar();
18   if (ch=='-') q=-1,ch=getchar();
19   while (ch>='0' && ch<='9') x=x*10+ch-'0',ch=getchar();
20   return q*x;
21 }
22 
23 il void insert(RG int from,RG int to){
24   g[++num]=(edge){head[from],to},head[from]=num; return;
25 }
26 
27 il void getrt(RG int x,RG int p,RG int &rt){
28   son[x]=0,sz[x]=1;
29   for (RG int i=head[x],v;i;i=g[i].nt){
30     v=g[i].to; if (v==p || vis[v]) continue;
31     getrt(v,x,rt),sz[x]+=sz[v],son[x]=max(son[x],sz[v]);
32   }
33   son[x]=max(son[x],son[0]-sz[x]);
34   if (son[rt]>=son[x]) rt=x; return;
35 }
36 
37 il void getdis(RG int x,RG int p){
38   dis[x]=dis[p]+a[x],st[++top]=(data){x,dis[x]},sz[x]=1;
39   for (RG int i=head[x],v;i;i=g[i].nt){
40     v=g[i].to; if (v==p || vis[v]) continue;
41     getdis(v,x),sz[x]+=sz[v];
42   }
43   return;
44 }
45 
46 il void calc(RG int rt,RG int p,RG int fg){
47   top=0,getdis(rt,p);
48   for (RG int i=1;i<=top;++i) ++tong[n+st[i].l];
49   for (RG int i=1,res;i<=top;++i)
50     res=fg*tong[n+(p?a[p]:a[rt])-st[i].l],sum+=res,f[st[i].i]+=res;
51   for (RG int i=1;i<=top;++i) --tong[n+st[i].l]; return;
52 }
53 
54 il void solve(RG int x,RG int S){
55   RG int rt=0; son[0]=S,getrt(x,0,rt);
56   vis[rt]=1,dis[rt]=a[rt],calc(rt,0,1);
57   for (RG int i=head[rt];i;i=g[i].nt)
58     if (!vis[g[i].to]) calc(g[i].to,rt,-1);
59   for (RG int i=head[rt];i;i=g[i].nt)
60     if (!vis[g[i].to]) solve(g[i].to,sz[g[i].to]);
61   return;
62 }
63 
64 int main(){
65 #ifndef ONLINE_JUDGE
66   freopen("king.in","r",stdin);
67   freopen("king.out","w",stdout);
68 #endif
69   n=gi();
70   for (RG int i=1;i<=n;++i) a[i]=gi()?1:-1;
71   for (RG int i=1,u,v;i<n;++i)
72     u=gi(),v=gi(),insert(u,v),insert(v,u);
73   solve(1,n);
74   for (RG int i=1,j=1;i<=n;++i){
75     now+=f[i]; while (j<=i && now<<1>sum) now-=f[j++];
76     ans+=j-1;
77   }
78   cout<<ans; return 0;
79 }