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

神仙啊，完全想不到……
又学了种套路（

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

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

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

代码：

#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));
    }
}