JZOJ-senior-5926. 【NOIP2018模拟10.25】naive 的图

Time Limits: 2000 ms Memory Limits: 524288 KB

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

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

Solution

显然，只有在最小生成树上的边才会对答案有贡献
我们对于每个点开一棵线段树，记录它所在集合里每种颜色出现的次数
在合并两个集合的时候，我们将 $size$ 较小的那个集合里的数在另一个集合里查询可以符合要求 $|c[u]-c[v]|>=L$ 的点的个数，总和记为 $sum$
边权 $\ast$ sum 即为这条边的贡献值
每条边的贡献值相加即为所求答案

Code

#include<algorithm>
#include<cstdio>

#define fo(i,a,b) for(int i=a;i<=b;++i)
#define fd(i,a,b) for(int i=a;i>=b;--i)
#define ll long long

using namespace std;

const int N=2e5+5,M=5e5+5,MX=1e9;
int n,m,L,tot,num;
int c[N],f[N],s[N],last[N];
ll sum,ans=0;
struct edge{int x,y,z;}e[M];
struct edge1{int to,next;}e1[N];
struct tree{int l,r,s;}tr[100*N];

bool cmp(edge x,edge y)
{
	return x.z<y.z;
}

int get(int x)
{
	return f[x]=f[x]==x?x:get(f[x]);	
}

void link(int x,int y)
{
	e1[++num]=(edge1){y,last[x]},last[x]=num;
}

void update(int x)
{
	int ls=tr[x].l,rs=tr[x].r;
	tr[x].s=tr[ls].s+tr[rs].s;
}

void ins(int x,int l,int r,int p)
{
	if(l==r) {++tr[x].s; return;}
	int mid=(l+r)>>1;
	if(p<=mid)
	{
		if(!tr[x].l) tr[x].l=++tot;
		ins(tr[x].l,l,mid,p);
	}
	else
	{
		if(!tr[x].r) tr[x].r=++tot;
		ins(tr[x].r,mid+1,r,p);
	}
	update(x);
}

int ask(int x,int st,int en,int l,int r)
{
	if(l>r) return 0;
	if(st==l&&en==r) return tr[x].s;
	int mid=(st+en)>>1;
	if(r<=mid) return ask(tr[x].l,st,mid,l,r);
	else if(l>mid) return ask(tr[x].r,mid+1,en,l,r);
	else return ask(tr[x].l,st,mid,l,mid)+ask(tr[x].r,mid+1,en,mid+1,r);
}

void merge(int x,int y,int l,int r)
{
	if(l==r) {tr[y].s+=tr[x].s; return;}
	ll mid=(l+r)>>1;
	if(tr[x].l&&tr[y].l) merge(tr[x].l,tr[y].l,l,mid);
		else if(tr[x].l) tr[y].l=tr[x].l;
	if(tr[x].r&&tr[y].r) merge(tr[x].r,tr[y].r,mid+1,r);
		else if(tr[x].r) tr[y].r=tr[x].r;
	update(y);
}

void dfs(int rt,int x,int fa)
{
	sum+=(ll)ask(rt,0,MX,0,c[x]-L);
	sum+=(ll)ask(rt,0,MX,c[x]+L,MX);
	for(int w=last[x];w;w=e1[w].next)
		if(e1[w].to!=fa) dfs(rt,e1[w].to,x);
}

int main()
{
	freopen("graph.in","r",stdin);
	freopen("graph.out","w",stdout);
	scanf("%d%d%d",&n,&m,&L),tot=n;
	fo(i,1,n)
	{
		scanf("%d",&c[i]);
		f[i]=i,s[i]=1;
		ins(i,0,MX,c[i]);
	}
	fo(i,1,m) scanf("%d%d%d",&e[i].x,&e[i].y,&e[i].z);
	sort(e+1,e+1+m,cmp);
	int cnt=0,i=1;
	while(cnt<n-1)
	{
		int fx=get(e[i].x),fy=get(e[i].y);
		if(fx!=fy)
		{
			if(s[fx]>s[fy]) swap(e[i].x,e[i].y),swap(fx,fy);
			if(!L) sum=(ll)s[fx]*(ll)s[fy];
				else sum=0,dfs(fy,fx,0);
			ans+=(ll)sum*(ll)e[i].z;
			link(fy,fx);
			merge(fx,fy,0,MX);
			f[fx]=fy,s[fy]+=s[fx],s[fx]=0;
			++cnt;
		}
		++i;
	}
	printf("%lld",ans);
}