JZOJ6866. 【2020.11.16提高组模拟】路径大小差

Description

• $n,w_i\le 2e5$

Solution

• 一道简单题，比赛的时候没有时间想。。。
• 可以简单点分治，合并讨论max,min的大小关系，用一个树状数组统计。
• 也可以容斥，求恰好$[l,l+k]$，相当于$[l,l+k]-[l+1,l+k]-[l,l+k-1]+[l+1,l+k-1]$，容斥一下，统计边权在$[l,l+m]$内的路径条数，可以线段树分治+按秩合并并查集支持撤销即可。

#include<cstdio>
#include<cmath>
#include<algorithm>
#include<cstring>
#include<vector>
#define maxn 200005
#define ll long long 
using namespace std;

int n,m,i,j,k,mx,mi,a[maxn][3],len;
ll ans,sum;

vector<int> t[maxn*4],bz[maxn*4];
int fa[maxn],sz[maxn];

void maketree(int x,int l,int r){
	t[x].clear(),bz[x].clear();
	if (l==r) return;
	int mid=(l+r)>>1;
	maketree(x<<1,l,mid),maketree(x<<1^1,mid+1,r);
}

void add(int x,int l,int r,int L,int R,int p){
	if (l>R||r<L) return;
	if (L<=l&&r<=R) {t[x].push_back(p),bz[x].push_back(0);return;}
	int mid=(l+r)>>1;
	add(x<<1,l,mid,L,R,p),add(x<<1^1,mid+1,r,L,R,p);
}

int father(int x){while (fa[x]) x=fa[x];return x;}

void merge(int x,int l,int r){
	for(int i=0;i<t[x].size();i++){
		int u=a[t[x][i]][0],v=a[t[x][i]][1];
		u=father(u),v=father(v),sum+=1ll*sz[u]*sz[v];
		if (sz[u]>sz[v]) fa[v]=u,sz[u]+=sz[v]; 
			else fa[u]=v,sz[v]+=sz[u],bz[x][i]=1;
	}
	if (l==r) {
		if (len==m) ans+=sum; else 
		if (len==m-2) ans+=sum*(l>1&&l+len<mx);else 
			ans-=sum+sum*(l>1&&l+len<mx);
	} else {
		int mid=(l+r)>>1;
		merge(x<<1,l,mid),merge(x<<1^1,mid+1,r);
	}
	for(int i=t[x].size()-1;i>=0;i--){
		int u=a[t[x][i]][0],v=a[t[x][i]][1],w=father(u),p;
		if (bz[x][i]) p=u; else p=v;
		while (fa[p]!=w) p=fa[p];
		sum-=1ll*(sz[w]-sz[p])*sz[p];
		sz[w]-=sz[p],fa[p]=0;
	}
}

void doit(){
	if (m<0||m>mx-mi) return;
	for(i=1;i<=n;i++) fa[i]=0,sz[i]=1;
	maketree(1,1,mx-len);
	memset(t,0,sizeof(t)),sum=0;
	for(i=1;i<n;i++) add(1,1,mx-len,a[i][2]-len,a[i][2],i);
	merge(1,1,mx-len);
}

int main(){
//	freopen("ceshi.in","r",stdin);
	freopen("minmax.in","r",stdin);
	freopen("minmax.out","w",stdout);
	scanf("%d%d",&n,&m),mx=0,mi=2e9;
	for(i=1;i<n;i++) scanf("%d%d%d",&a[i][0],&a[i][1],&a[i][2]);
	for(i=1;i<n;i++) mx=max(mx,a[i][2]),mi=min(mi,a[i][2]);
	len=m,doit();
	len=m-1,doit();
	len=m-2,doit();
	printf("%lld\n",ans);
}