[POJ 1741] Tree

一、題目

點此看題

題目描述

給一個$n$個點的數，找出多少點對$(u,v)$（$u<v$）的距離不大於$k$

二、解法

可以用啓發式合併，我寫的是無旋$treap$，康康代碼吧。

#include <cstdio>
#include <cstdlib>
#include <iostream>
using namespace std;
const int M = 100005;
int read()
{
    int num=0,flag=1;
    char c;
    while((c=getchar())<'0'||c>'9')if(c=='-')flag=-1;
    while(c>='0'&&c<='9')num=(num<<3)+(num<<1)+(c^48),c=getchar();
    return num*flag;
}
int n,k,tot,ans,cnt,f[M],rt[M],ch[M][2],siz[M],val[M],hp[M],la[M];
struct node
{
	int p[2];
	node() {p[0]=p[1]=0;}
}emp;
struct edge
{
	int v,c,next;
}e[2*M];
void up(int x)
{
	if(!x) return ;
	siz[x]=siz[ch[x][0]]+siz[ch[x][1]]+1;
}
void change(int x,int v)
{
	if(!x) return ;
	la[x]+=v;
	val[x]+=v;
}
void down(int x)
{
	if(!la[x] || !x) return ;
	change(ch[x][0],la[x]);
	change(ch[x][1],la[x]);
	la[x]=0;
}
node split(int x,int v)
{
	if(!x) return emp;
	int d=val[x]<=v;
	down(x);
	node y=split(ch[x][d],v);
	ch[x][d]=y.p[d^1];
	y.p[d^1]=x;
	up(x);
	return y;
}
int merge(int x,int y)
{
	if(!x || !y) return x+y;
	if(hp[x]<hp[y])
	{
		down(x);
		ch[x][1]=merge(ch[x][1],y);
		up(x);
		return x;
	}
	down(y);
	ch[y][0]=merge(x,ch[y][0]);
	up(y);
	return y;
}
void fuck(int x,int y)
{
	if(!x) return ;
	down(x);
	fuck(ch[x][0],y);
	fuck(ch[x][1],y);
	if(val[x]<=k)
	{
		node t=split(y,k-val[x]);
		ans+=siz[t.p[0]];
		y=merge(t.p[0],t.p[1]);
	}
}
void uni(int x,int &y)
{
	if(!x) return ;
	uni(ch[x][0],y);
	uni(ch[x][1],y);
	ch[x][0]=ch[x][1]=0;siz[x]=1;
	node t=split(y,val[x]);
	y=merge(t.p[0],merge(x,t.p[1]));
}
void dfs(int u,int fa)
{
	rt[u]=++cnt;siz[cnt]=1;hp[cnt]=rand();
	for(int i=f[u];i;i=e[i].next)
	{
		int v=e[i].v,c=e[i].c;
		if(v==fa) continue;
		dfs(v,u);
		change(rt[v],c);
		if(siz[rt[u]]<siz[rt[v]]) swap(rt[u],rt[v]);//啓發式合併 
		fuck(rt[v],rt[u]);//先統計答案 
		uni(rt[v],rt[u]);//暴力合併 
	}
}
signed main()
{
	while(~scanf("%d %d",&n,&k) && n && k)
	{
		cnt=tot=ans=0;
		for(int i=1;i<=n;i++)
		{
			f[i]=siz[i]=la[i]=0;
			ch[i][0]=ch[i][1]=val[i]=0;
		}
		for(int i=1;i<n;i++)
		{
			int u=read(),v=read(),c=read();
			e[++tot]=edge{v,c,f[u]},f[u]=tot;
			e[++tot]=edge{u,c,f[v]},f[v]=tot;
		}
		dfs(1,0);
		printf("%d\n",ans);
	}
}