LibreOJ 3227.「USACO 2019 · December」Bessie's Snow Cow

神仙啊,完全想不到……
又学了种套路(

有一个神奇的思路:对于每种颜色,维护所有被泼过此颜色的子树,但是若子树之间有包含关系则保留最大的
查询时应考虑祖先中是否泼过,以及子树内是否泼过。

可以对每个颜色开一个 set,并用两只 BIT 来维护答案:第一只维护祖先的,第二只维护子树的。
修改时,首先看父亲中有没有泼过同颜色的,如果有则忽略此次操作。
然后看子树内有没有泼过同颜色的,如果有则撤销贡献。
最后在 BIT 上增加这个修改的贡献:在第一只 BIT 中对子树全部加一,在第二只 BIT 中对该点加上子树大小

则查询时的答案为第一只 BIT 中该点的贡献乘上子树大小与第二只 BIT 中的子树和之和。
不过可能会重复计算,可以在统计第二只 BIT 中贡献时忽略该点,让第一只 BIT 去统计。

代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
#include <cstdio>
#include <set>
#define lowbit(x) ((x) & -(x))
using namespace std;
const int N = 1e5;
int n,m;
int to[(N << 1) + 5],pre[(N << 1) + 5],first[N + 5];
inline void add(int u,int v)
{
static int tot = 0;
to[++tot] = v,pre[tot] = first[u],first[u] = tot;
}
int fa[N + 5],rk[N + 5],id[N + 5],sz[N + 5];
void dfs(int p)
{
static int tot = 0;
rk[id[p] = ++tot] = p,sz[p] = 1;
for(register int i = first[p];i;i = pre[i])
if(to[i] ^ fa[p])
fa[to[i]] = p,dfs(to[i]),sz[p] += sz[to[i]];
}
long long c[2][N + 5];
void update(int op,int x,int k)
{
for(;x <= n;x += lowbit(x))
c[op][x] += k;
}
long long query(int op,int x)
{
long long ret = 0;
for(;x;x -= lowbit(x))
ret += c[op][x];
return ret;
}
void update(int p,int k)
{
update(0,id[p],k),update(0,id[p] + sz[p],-k),update(1,id[p],k * sz[p]);
}
long long query(int p)
{
return query(0,id[p]) * sz[p] + query(1,id[p] + sz[p] - 1) - query(1,id[p]);
}
set<int> s[N + 5];
int main()
{
scanf("%d%d",&n,&m);
int u,v;
for(register int i = 1;i < n;++i)
scanf("%d%d",&u,&v),add(u,v),add(v,u);
dfs(1);
for(register int i = 1;i <= N;++i)
s[i].insert(0),s[i].insert(n + 1);
for(int op,x,c,p;m;--m)
{
scanf("%d%d",&op,&x);
if(op == 1)
{
scanf("%d",&c),p = rk[*--s[c].upper_bound(id[x])];
if(id[x] >= id[p] + sz[p])
{
for(register set<int>::iterator it = s[c].upper_bound(id[x]);*it < id[x] + sz[x];update(rk[*it],-1),s[c].erase(it),it = s[c].upper_bound(id[x]));
update(x,1),s[c].insert(id[x]);
}
}
else
printf("%lld\n",query(x));
}
}