【JZOJ A组】naive 的图

Description

众所周知，小 naive 有一张 n 个点，m 条边的带权无向图。第 i 个点的颜色为 ci。d(s, t)表示从点 s 到点 t 的权值最小的路径的权值，一条路径的权值定义为路径上权值最大的边的权值。
求所有满足 u < v, |cu − cv| ≥ L 的点对 (u, v) 的 d(u, v) 之和。

Input

输入文件为 graph.in。
第一行，三个整数 n, m, L，表示点数，边数和参数 L。
第二行，n 个整数，第 i 个数为第 i 个点的颜色 ci。接下来 m 行，每行三个整数 ui, vi, wi，描述了一条边。

Output

输出文件为 graph.out。
共一行，一个整数，表示答案。

Sample Input

4 5 2
6 4 5 2
1 2 8
2 3 7
2 4 8
1 3 2
1 4 1

Sample Output

17

样例解释
满足条件的点对：(1, 2),(1, 4),(2, 4),(3, 4)，答案为 7 + 1 + 7 + 2 = 17。

Data Constraint

对于每个测试点内的数据：

思路

先建出 Kruskal 重构树，每条边的贡献次数为它连接的两个子树之间的颜色之差大于等于 L 的点对数，可以发现 ∑ min(size(lef tchildi), size(rightchildi)) = O(nlog2n)。
对于每条边我们枚举 size 较小的那棵子树内的点，算出在另一棵子树中能与它组成点对的点的个数。这个问题实际上就是询问在 dfs 序的一段区间上并且颜色不在一段区间内的点数，二维数点问题可以离线树状数组完成。
总的时间复杂度为 O(mlog2m + nlog2n)

代码

#include<cstdio>
#include<iostream>
#include<algorithm>
#define ll long long
using namespace std;
const int inf=0x3f3f3f3f,N=2e5+77,M=5e5+77;
ll n,m,l,tot,num,sum,ans,c[N],f[N],d[N],list[N];
struct E
{
	ll x,y,z;
}a[M];
struct node
{
	ll to,next;
}e[N];
struct tr{ ll l,r,v; }tr[100*N];
bool cmp(E a,E b) { return a.z<b.z; }
ll gf(ll x) { return f[x]==x?x:f[x]=gf(f[x]); }
void add(ll x,ll y) { e[++tot].to=y; e[tot].next=list[x]; list[x]=tot; }
ll query(ll d,ll l,ll r,ll st,ll ed)
{
	if(st>ed) return 0;
	if(l==st&&r==ed) return tr[d].v;
	ll mid=(l+r)>>1;
	if(ed<=mid) return query(tr[d].l,l,mid,st,ed);
	else if(st>mid) return query(tr[d].r,mid+1,r,st,ed);
	else return query(tr[d].l,l,mid,st,mid)+query(tr[d].r,mid+1,r,mid+1,ed);
}
void dfs(ll d,ll x,ll fa)
{
	sum+=query(d,0,inf,0,c[x]-l)+query(d,0,inf,c[x]+l,inf);
	for(ll i=list[x]; i; i=e[i].next) if(e[i].to!=fa) dfs(d,e[i].to,x);
}
void ins(ll d,ll l,ll r,ll x)
{
	if(l==r) { tr[d].v++; return; }
	ll mid=(l+r)>>1;
	if(x<=mid)
	{
		if(!tr[d].l) tr[d].l=++num;
		ins(tr[d].l,l,mid,x);
	}
	else 
	{
		if(!tr[d].r) tr[d].r=++num;
		ins(tr[d].r,mid+1,r,x);
	}
	tr[d].v=tr[tr[d].l].v+tr[tr[d].r].v;
}
void merge(ll l,ll r,ll st,ll ed)
{
	if(st==ed) 
	{
		tr[r].v+=tr[l].v; return;
	}
	ll mid=(st+ed)>>1;
	if(tr[l].l&&tr[r].l) merge(tr[l].l,tr[r].l,st,mid); else if(tr[l].l) tr[r].l=tr[l].l;
	if(tr[l].r&&tr[r].r) merge(tr[l].r,tr[r].r,mid+1,ed); else if(tr[l].r) tr[r].r=tr[l].r;
	tr[r].v=tr[tr[r].l].v+tr[tr[r].r].v;
}
int main()
{
	freopen("graph.in","r",stdin),freopen("graph.out","w",stdout);
	scanf("%lld%lld%lld",&n,&m,&l),num=n;
	for(ll i=1; i<=n; i++) scanf("%lld",&c[i]),f[i]=i,d[i]=1,ins(i,0,inf,c[i]);
	for(ll i=1; i<=m; i++) scanf("%lld%lld%lld",&a[i].x,&a[i].y,&a[i].z);
	sort(a+1,a+m+1,cmp);
	ll i=1,tot=0;
	while(tot<n-1)
	{
		ll u=gf(a[i].x),v=gf(a[i].y);
		if(u!=v)
		{
			if(d[u]>d[v]) swap(a[i].x,a[i].y),swap(u,v);
			if(!l) sum=d[u]*d[v]; else sum=0,dfs(v,u,0);
			ans+=sum*a[i].z,add(v,u),merge(u,v,0,inf);
			f[u]=v,d[v]+=d[u];d[u]=0;
			++tot;
		}
		i++;
	}
	printf("%lld\n",ans);
}