【洛谷P4074】糖果公园

题目

题目链接：https://www.luogu.com.cn/problem/P4074
Candyland 有一座糖果公园，公园里不仅有美丽的风景、好玩的游乐项目，还有许多免费糖果的发放点，这引来了许多贪吃的小朋友来糖果公园游玩。
糖果公园的结构十分奇特，它由 \(n\) 个游览点构成，每个游览点都有一个糖果发放处，我们可以依次将游览点编号为 \(1\) 至 \(n\)。有 \(n\) – \(1\) 条双向道路连接着这些游览点，并且整个糖果公园都是连通的，即从任何一个游览点出发都可以通过这些道路到达公园里的所有其它游览点。
糖果公园所发放的糖果种类非常丰富，总共有 \(m\) 种，它们的编号依次为 \(1\) 至 \(m\)。每一个糖果发放处都只发放某种特定的糖果，我们用 \(C_i\) 来表示 \(i\) 号游览点的糖果。
来到公园里游玩的游客都不喜欢走回头路，他们总是从某个特定的游览点出发前往另一个特定的游览点，并游览途中的景点，这条路线一定是唯一的。他们经过每个游览点，都可以品尝到一颗对应种类的糖果。
大家对不同类型糖果的喜爱程度都不尽相同。 根据游客们的反馈打分，我们得到了糖果的美味指数， 第 \(i\) 种糖果的美味指数为 \(V_i\)。另外，如果一位游客反复地品尝同一种类的糖果，他肯定会觉得有一些腻。根据量化统计，我们得到了游客第 \(i\) 次品尝某类糖果的新奇指数 \(W_i\)。如果一位游客第 \(i\) 次品尝第 \(j\) 种糖果，那么他的愉悦指数 \(H\) 将会增加对应的美味指数与新奇指数的乘积，即 \(V_j\)×\(W_i\)。这位游客游览公园的愉悦指数最终将是这些乘积的和。
当然，公园中每个糖果发放点所发放的糖果种类不一定是一成不变的。有时，一些糖果点所发放的糖果种类可能会更改（也只会是 \(m\) 种中的一种），这样的目的是能够让游客们总是感受到惊喜。
糖果公园的工作人员小 A 接到了一个任务，那就是根据公园最近的数据统计出每位游客游玩公园的愉悦指数。但数学不好的小 A 一看到密密麻麻的数字就觉得头晕，作为小 A 最好的朋友，你决定帮他一把。

思路

好久没敲莫队了，于是敲一个基础莫队合集练练手（
树上莫队一般维护括号序，区间 \([l,r]\) 内一个点有贡献当且仅当该点在 \([l,r]\) 中只出现了一次。试想一下，将 \(x\) 的子树遍历完毕并回溯时，\(x\) 就会被遍历第二遍，且此时 \(x\) 就不在路径中了。
所以打一个标记记录每个点在区间内出现次数即可。再套上时间戳，解决修改。
\(T\) 取 \(n^{\frac{2}{3}}\) 时复杂度为 \(O(n^{\frac{5}{3}})\)。
代码奇丑。因为众所周知莫队合并语句会快很多。

代码

#pragma GCC optimize("Ofast","inline")
#include <cmath>
#include <cstdio>
#include <cctype>
#include <cstring>
#include <algorithm>
#define reg register
using namespace std;
typedef long long ll;

const int N=200010,LG=20;
int head[N],pos[N],st[N],nd[N],cnt[N],v[N],w[N],c[N],a[N],f[N][LG+1],dep[N];
int n,m,Q,Q1,Q2,T,tot;
bool flag[N];
ll ans,Ans[N];

inline int read()
{
	int d=0; char ch=getchar();
	while (!isdigit(ch)) ch=getchar();
	while (isdigit(ch)) d=(d<<3)+(d<<1)+ch-48,ch=getchar();
	return d;
}

struct Query
{
	int l,r,t,num;
}ask[N];

inline bool operator < (const Query &a,const Query &b)
{
	return pos[a.l]^pos[b.l] ? a.l<b.l : pos[a.l]&1 ? a.r<b.r : a.r>b.r;
}

struct Change
{
	int x,y;
}chg[N];

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

inline bool cmp(Query x,Query y)
{
	if (pos[x.l]<pos[y.l]) return 1;
	if (pos[x.l]>pos[y.l]) return 0;
	return x.r<y.r;
}

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

inline void dfs(int x,int fa)
{
	a[++tot]=x; st[x]=tot; pos[tot]=tot/T;
	f[x][0]=fa; dep[x]=dep[fa]+1;
	for (reg int i=1;i<=LG;i++)
		f[x][i]=f[f[x][i-1]][i-1];
	for (reg int i=head[x];~i;i=e[i].next)
		if (e[i].to!=fa) dfs(e[i].to,x);
	a[++tot]=x; nd[x]=tot; pos[tot]=tot/T;
}

inline int lca(int x,int y)
{
	if (dep[x]<dep[y]) swap(x,y);
	for (reg int i=LG;i>=0;i--)
		if (dep[f[x][i]]>=dep[y]) x=f[x][i];
	if (x==y) return x;
	for (reg 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];
}

inline void ins(int x)
{
	ans+=1LL*v[c[x]]*w[++cnt[c[x]]];
	flag[x]^=1;
}

inline void del(int x)
{
	ans-=1LL*v[c[x]]*w[cnt[c[x]]--];
	flag[x]^=1;
}

inline void solve(int x)
{
	(flag[x]^=1) ? ans+=1LL*v[c[x]]*w[++cnt[c[x]]] : ans-=1LL*v[c[x]]*w[cnt[c[x]]--];
}

inline void update(int i)
{
	if (flag[chg[i].x])
	{
		int y=chg[i].y,col=c[chg[i].x];
		ans=ans-1LL*v[col]*w[cnt[col]--]+1LL*v[y]*w[++cnt[y]];
	}
	swap(chg[i].y,c[chg[i].x]);
}

int main()
{
	memset(head,-1,sizeof(head));
	n=read(); m=read(); Q=read();
	T=pow(n,2.0/3);
	for (reg int i=1;i<=m;i++) v[i]=read();
	for (reg int i=1;i<=n;i++) w[i]=read();
	for (reg int i=1,x,y;i<n;i++)
	{
		x=read(); y=read();
		add(x,y); add(y,x);
	}
	for (reg int i=1;i<=n;i++) c[i]=read();
	tot=0; dfs(1,0);
	for (reg int i=1,type;i<=Q;i++)
	{
		type=read();
		if (!type)
		{
			Q2++;
			chg[Q2].x=read(); chg[Q2].y=read();
		}
		else 
		{
			Q1++;
			ask[Q1].t=Q2; ask[Q1].num=Q1;
			int x=read(),y=read(),p=lca(x,y);
			if (st[x]>st[y]) swap(x,y);
			if (p==x || p==y)
				ask[Q1].l=st[x],ask[Q1].r=st[y];
			else
				ask[Q1].l=nd[x],ask[Q1].r=st[y];
		}
	}
	sort(ask+1,ask+1+Q1);
	for (reg int l=1,r=0,t=0,i=1;i<=Q1;i++)
	{		
		while (l<ask[i].l) solve(a[l++]);
		while (l>ask[i].l) solve(a[--l]);
		while (r<ask[i].r) solve(a[++r]);
		while (r>ask[i].r) solve(a[r--]);
		
		for (;t<ask[i].t;t++) update(t+1);
		for (;t>ask[i].t;t--) update(t);
		int p=lca(a[l],a[r]);
		if (a[l]!=p && a[r]!=p) ins(p);
		Ans[ask[i].num]=ans;
		if (a[l]!=p && a[r]!=p) del(p);
	}
	for (reg int i=1;i<=Q1;i++)
		printf("%lld\n",Ans[i]);
	return 0;
}