【洛谷P3320】寻宝游戏

题目

题目链接：https://www.luogu.com.cn/problem/P3320
小B最近正在玩一个寻宝游戏，这个游戏的地图中有N个村庄和N-1条道路，并且任何两个村庄之间有且仅有一条路径可达。游戏开始时，玩家可以任意选择一个村庄，瞬间转移到这个村庄，然后可以任意在地图的道路上行走，若走到某个村庄中有宝物，则视为找到该村庄内的宝物，直到找到所有宝物并返回到最初转移到的村庄为止。
小B希望评测一下这个游戏的难度，因此他需要知道玩家找到所有宝物需要行走的最短路程。但是这个游戏中宝物经常变化，有时某个村庄中会突然出现宝物，有时某个村庄内的宝物会突然消失，因此小B需要不断地更新数据，但是小B太懒了，不愿意自己计算，因此他向你求助。为了简化问题，我们认为最开始时所有村庄内均没有宝物

思路

假设1为树根，求出每一个点的\(dfs\)序，那么要便利现在所有的宝藏且路径最短，只需按照\(dfs\)序排序，然后依次走一遍即可。
所以我们需要一个支持插入、求前驱后继、删除的数据结构，为了满足\(dfs\)序最小的要和\(dfs\)序最大的也连在一起，又需要支持查询最值。这明显是平衡树的板子，直接大力上\(Treap\)即可。
当然，权值线段树和二分+树状数组也可以达到同样效果。
时间复杂度\(O(m\log n)\)。

代码

别问我为什么打了两份。。。我一开始写Treap以为写炸了，本地全对，交上去\(0pts\)。于是改写bit+二分，后来写完之后发现原来那份忘记关\(freopen\)了。。。

Treap

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;

const int N=100010,LG=20,Inf=1e9;
int n,m,tot,root,head[N],dfn[N],rk[N],f[N][LG+1],dep[N];
ll ans,dis[N];
bool flag[N];

struct edge
{
	int next,to,dis;
}e[N*2];

void add(int from,int to,int dis)
{
	e[++tot].to=to;
	e[tot].dis=dis;
	e[tot].next=head[from];
	head[from]=tot;
}

void dfs(int x,int fa)
{
	f[x][0]=fa; dep[x]=dep[fa]+1;
	dfn[x]=++tot; rk[tot]=x;
	for (int i=1;i<=LG;i++)
		f[x][i]=f[f[x][i-1]][i-1];
	for (int i=head[x];~i;i=e[i].next)
		if (e[i].to!=fa)
		{
			dis[e[i].to]=dis[x]+e[i].dis;
			dfs(e[i].to,x);
		}
}

int lca(int x,int y)
{
	if (dep[x]<dep[y]) swap(x,y);
	for (int i=LG;i>=0;i--)
		if (dep[f[x][i]]>=dep[y]) x=f[x][i];
	if (x==y) return x;
	for (int i=LG;i>=0;i--)
		if (f[x][i]!=f[y][i]) x=f[x][i],y=f[y][i];
	return f[x][0];
}

struct Treenode
{
	int lc,rc,cnt,val,size,dat;
};

struct Treap
{
	Treenode t[N];
	int tot;
	
	int New(int val)
	{
		t[++tot].val=val;
		t[tot].dat=rand();
		t[tot].cnt=t[tot].size=1;
		return tot;
	}
	
	void update(int x)
	{
		t[x].size=t[t[x].lc].size+t[t[x].rc].size+t[x].cnt;
	}
	
	void build()
	{
		root=New(-Inf);
		t[1].rc=New(Inf);
		update(1);
	}
	
	void zig(int &x)
	{
		int y=t[x].lc;
		t[x].lc=t[y].rc; t[y].rc=x; x=y;
		update(t[x].rc); update(x);
	}
	
	void zag(int &x)
	{
		int y=t[x].rc;
		t[x].rc=t[y].lc; t[y].lc=x; x=y;
		update(t[x].lc); update(x);
	}
	
	void insert(int &x,int val)
	{
		if (!x)
		{
			x=New(val);
			return;
		}
		if (t[x].val==val)
		{
			t[x].cnt++;
			update(x);
			return;
		}
		if (val<t[x].val)
		{
			insert(t[x].lc,val);
			if (t[x].dat<t[t[x].lc].dat) zig(x);
		}
		else
		{
			insert(t[x].rc,val);
			if (t[x].dat<t[t[x].rc].dat) zag(x);
		}
		update(x);
	}
	
	void del(int &x,int val)
	{
		if (!x) return;
		if (t[x].val==val)
		{
			if (t[x].cnt>1)
			{
				t[x].cnt--;
				update(x);
				return;
			}
			if (t[x].lc || t[x].rc)
			{
				if (!t[x].lc || t[t[x].rc].dat>t[t[x].lc].dat)
					zag(x),del(t[x].lc,val);
				else
					zig(x),del(t[x].rc,val);
				update(x);
			}
			else x=0;
			return;
		}
		if (val<t[x].val) del(t[x].lc,val);
			else del(t[x].rc,val);
		update(x);
	}
	
	int get_rank(int x,int val)
	{
		if (!x) return 0;
		if (val==t[x].val) return t[t[x].lc].size+1;
		if (val<t[x].val) return get_rank(t[x].lc,val);
			else return get_rank(t[x].rc,val)+t[t[x].lc].size+t[x].cnt;
	}
	
	int get_val(int x,int rank)
	{
		if (!x) return Inf;
		if (t[t[x].lc].size+1<=rank && t[t[x].lc].size+t[x].cnt>=rank)
			return t[x].val;
		if (rank<=t[t[x].lc].size) return get_val(t[x].lc,rank);
			else return get_val(t[x].rc,rank-t[t[x].lc].size-t[x].cnt);
	}
	
	int pre(int x,int val)
	{
		if (!x) return -Inf;
		if (t[x].val<val) return max(t[x].val,pre(t[x].rc,val));
			else return pre(t[x].lc,val);
	}
	
	int nxt(int x,int val)
	{
		if (!x) return Inf;
		if (t[x].val>val) return min(t[x].val,nxt(t[x].lc,val));
			else return nxt(t[x].rc,val);
	}
	
	int Min(int x)
    {
        if (!x) return Inf;
        if (t[x].val==-Inf) return Min(t[x].rc);
        return min(Min(t[x].lc),t[x].val);
    }
    
    int Max(int x)
    {
        if (!x) return -Inf;
        if (t[x].val==Inf) return Max(t[x].lc);
        return max(Max(t[x].rc),t[x].val);
    }
}Treap;

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&m);
	for (int i=1,x,y,z;i<n;i++)
	{
		scanf("%d%d%d",&x,&y,&z);
		add(x,y,z); add(y,x,z);
	}
	tot=0; dfs(1,0);
	while (m--)
	{
		int x;
		scanf("%d",&x);
		if (flag[x])
		{
			int p=Treap.pre(root,dfn[x]),q=Treap.nxt(root,dfn[x]);
			if (p==-Inf) p=Treap.Max(root);
			if (q==Inf) q=Treap.Min(root);
			if (-Inf<p && p<Inf)
			{
				int LCA=lca(x,rk[p]);
				ans-=dis[x]-dis[LCA]+dis[rk[p]]-dis[LCA];
			}
			if (-Inf<q && q<Inf)
			{
				int LCA=lca(x,rk[q]);
				ans-=dis[x]-dis[LCA]+dis[rk[q]]-dis[LCA];
			}
			if (-Inf<p && p<Inf && -Inf<q && q<Inf)
			{
				int LCA=lca(rk[p],rk[q]);
				ans+=dis[rk[p]]-dis[LCA]+dis[rk[q]]-dis[LCA];
			}
			Treap.del(root,dfn[x]);
		}
		else
		{
			int p=Treap.pre(root,dfn[x]),q=Treap.nxt(root,dfn[x]);
			if (p==-Inf) p=Treap.Max(root);
			if (q==Inf) q=Treap.Min(root);
			if (-Inf<p && p<Inf)
			{
				int LCA=lca(x,rk[p]);
				ans+=dis[x]-dis[LCA]+dis[rk[p]]-dis[LCA];
			}
			if (-Inf<q && q<Inf)
			{
				int LCA=lca(x,rk[q]);
				ans+=dis[x]-dis[LCA]+dis[rk[q]]-dis[LCA];
			}
			if (-Inf<p && p<Inf && -Inf<q && q<Inf)
			{
				int LCA=lca(rk[p],rk[q]);
				ans-=dis[rk[p]]-dis[LCA]+dis[rk[q]]-dis[LCA];
			}
			Treap.insert(root,dfn[x]);
		}
		flag[x]^=1;
		printf("%lld\n",ans);
	}
	return 0;
}

树状数组+二分

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;

const int N=100010,LG=20,Inf=1e9;
int n,m,tot,cnt,head[N],dfn[N],rk[N],f[N][LG+1],dep[N];
ll ans,dis[N];
bool flag[N];

struct edge
{
	int next,to,dis;
}e[N*2];

struct BIT
{
	ll c[N];
	
	BIT()
	{
		memset(c,0,sizeof(c));
	}
	
	void add(int x,ll val)	
	{
		for (;x<=n;x+=x&-x)
			c[x]+=val;
	}
	
	int ask(int x)
	{
		ll sum=0;
		for (;x;x-=x&-x)
			sum+=c[x];
		return sum;
	}
}bit;

void add(int from,int to,int dis)
{
	e[++tot].to=to;
	e[tot].dis=dis;
	e[tot].next=head[from];
	head[from]=tot;
}

void dfs(int x,int fa)
{
	f[x][0]=fa; dep[x]=dep[fa]+1;
	dfn[x]=++tot; rk[tot]=x;
	for (int i=1;i<=LG;i++)
		f[x][i]=f[f[x][i-1]][i-1];
	for (int i=head[x];~i;i=e[i].next)
		if (e[i].to!=fa)
		{
			dis[e[i].to]=dis[x]+e[i].dis;
			dfs(e[i].to,x);
		}
}

int lca(int x,int y)
{
	if (dep[x]<dep[y]) swap(x,y);
	for (int i=LG;i>=0;i--)
		if (dep[f[x][i]]>=dep[y]) x=f[x][i];
	if (x==y) return x;
	for (int i=LG;i>=0;i--)
		if (f[x][i]!=f[y][i]) x=f[x][i],y=f[y][i];
	return f[x][0];
}

int binary(int k)
{
	int l=0,r=n;
	while (l<=r)
	{
		int mid=(l+r)>>1;
		if (bit.ask(mid)>=k) r=mid-1;
			else l=mid+1;
	}
	return r+1;
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&m);
	for (int i=1,x,y,z;i<n;i++)
	{
		scanf("%d%d%d",&x,&y,&z);
		add(x,y,z); add(y,x,z);
	}
	tot=0; dfs(1,0);
	while (m--)
	{
		int x;
		scanf("%d",&x);
		int pos=bit.ask(dfn[x]);
		if (!flag[x])
		{
			int p=binary(pos),q=binary(pos+1);
			if (!p || p>n) p=binary(cnt);
			if (!q || q>n) q=binary(1);
			if (1<=p && p<=n)
			{
				int LCA=lca(rk[p],x);
				ans+=dis[rk[p]]+dis[x]-2*dis[LCA];
			}
			if (1<=q && q<=n)
			{
				int LCA=lca(rk[q],x);
				ans+=dis[rk[q]]+dis[x]-2*dis[LCA];
			}
			if (1<=p && p<=n && 1<=q && q<=n)
			{
				int LCA=lca(rk[p],rk[q]);
				ans-=dis[rk[p]]+dis[rk[q]]-2*dis[LCA];
			}
			bit.add(dfn[x],1);
			cnt++;
		}
		else
		{
			int p=binary(pos-1),q=binary(pos+1);
			if (!p || p>n) p=binary(cnt);
			if (!q || q>n) q=binary(1);
			if (1<=p && p<=n)
			{
				int LCA=lca(rk[p],x);
				ans-=dis[rk[p]]+dis[x]-2*dis[LCA];
			}
			if (1<=q && q<=n)
			{
				int LCA=lca(rk[q],x);
				ans-=dis[rk[q]]+dis[x]-2*dis[LCA];
			}
			if (1<=p && p<=n && 1<=q && q<=n)
			{
				int LCA=lca(rk[p],rk[q]);
				ans+=dis[rk[p]]+dis[rk[q]]-2*dis[LCA];
			}
			bit.add(dfn[x],-1);
			cnt--;
		}
		flag[x]^=1;
		printf("%lld\n",ans);
	}
	return 0;
}