[USACO19DEC]Bessie's Snow Cow P 题解

题目链接： https://www.luogu.com.cn/problem/P5852

题意：

子树染色，子树询问每个点的颜色数之和，一个点可以有多种颜色。 \((1\le N,Q\le10^5)\)

题解：

首先来看看子树染一次色有什么贡献（假设这棵树还没染过这种颜色）。

• 对于这个子树内部的点，是所有点子树内的点都有了这种颜色，因此对这些点的贡献是子树大小。

• 其余点只有是祖先的才会有贡献，并且贡献都是这个点的子树大小。

这些东西可以开两个树状数组维护，一个查有多少祖先用颜色覆盖了他（第一种），再乘个子树大小，另一个直接在自己这里加上自己的子树大小，查询时查询整棵子树的贡献。

要是染过这种颜色怎么办呢？如果两点之间没有祖孙关系，那是没有影响的。如果有子孙关系，可以直接把子孙的贡献去掉。

正确性显然。而这个关系可以对每种颜色开一个 set 按 DFS序 维护，只有相邻的点会有祖孙关系，前面是后面的祖先。

这种东西可以直接按进和出的编号来弄……现在突然发现出点的时候不用加一，这样还可以顺便维护子树大小。另外，对于被修改的那个点两种贡献可能会重复，第二种不要算进去自身就好了……

时间复杂度： \(O(N\log N)\) 。

代码：

#include <bits/stdc++.h>
#define to e[x][i]
using namespace std;
typedef long long ll;
typedef set<int>::iterator SIT;
const int N=1e5+5;
int n,q,sum,dfi[N],dfo[N],rev[N];
vector<int> e[N];
set<int> s[N];
struct BIT{
	ll tree[N];
	inline int lb(int x) {return x&(-x);}
	inline void add(int x,int v) {while(x<=n) {tree[x]+=v;x+=lb(x);}}
	inline ll ask(int x) {ll res=0;while(x) {res+=tree[x];x-=lb(x);}return res;}
}t1,t2;
inline SIT pre(SIT x) {return --x;}
inline int upd(int x,int v) {t1.add(dfi[x],v);t1.add(dfo[x]+1,-v);t2.add(dfi[x],(dfo[x]-dfi[x]+1)*v);}
void dfs(int x,int fa){
	rev[dfi[x]=++sum]=x;
	for(int i=0;i<e[x].size();i++)
		if(to!=fa) dfs(to,x);
	dfo[x]=sum;
}
signed main(){
	scanf("%d%d",&n,&q);
	for(int i=1,u,v;i<n;i++) {scanf("%d%d",&u,&v);e[u].push_back(v);e[v].push_back(u);}
	dfs(1,0);
	while(q--){
		int op,x,y;scanf("%d%d",&op,&x);
		if(op==1){
			scanf("%d",&y);SIT it=s[y].upper_bound(dfi[x]);
			if(it!=s[y].begin()&&dfo[rev[*pre(it)]]>=dfo[x]) continue;
			while(it!=s[y].end()&&(*it)<=dfo[x]) {upd(rev[*it],-1);s[y].erase(it++);}
			s[y].insert(dfi[x]);upd(x,1);
		}
		else printf("%lld\n",t1.ask(dfi[x])*(dfo[x]-dfi[x]+1)+t2.ask(dfo[x])-t2.ask(dfi[x]));
	}
	return 0;
}