CodeForces 620E"New Year Tree"（DFS序+线段树+状态压缩）

传送门

•题意

　　给你一颗 n 个节点的树，每个节点被染上了颜色；

　　有 m 次操作，每次操作的类型有两种

• • 1 v c : 将以 v 为根的子树的结点全部涂成 c
  • 2 v : 询问以 v 为根的子树的结点中不同颜色的数量

•题解

　　因为操作的是某个节点对应的子树，所以我们可以通过DFS序获得每一个结点的子树区间。

　　之后我们就将这个树形的问题转化成了一个线性区间的问题；

　　然后，就可以用线段树来维护（区间更新，区间查询）；

　　求解某个区间不同颜色的个数要怎么办呢？

　　刚开始，我在线段树中定义了一个 $set$，存储当前区间所有的颜色；

　　但是，用 $set$ 超时了；

　　因为总共的颜色最多只有 60 种；

　　因此我们可以用一个范围在 $long long$ 内的二进制数 $bit$ 存储每一种颜色；

　　（$bit$ 的二进制下的每一位代表着一种颜色）

•Code

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define ls(x) (x<<1)
#define rs(x) (x<<1|1)
#define mem(a,b) memset(a,b,sizeof(a))
#define popcount(x) __builtin_popcountll(x)///获取x的二进制位1的个数
const int maxn=4e5+50;

int n,m;
int a[maxn];
int num;
int head[maxn];
struct Edge
{
    int to;
    int next;
}G[maxn<<1];
void addEdge(int u,int v)
{
    G[num]={v,head[u]};
    head[u]=num++;
}
struct Seg
{
    int l,r;
    ll sum;
    int lazy;
    int mid(){return l+((r-l)>>1);};
    void Set(int col)
    {
        sum=1ll<<col;
        lazy=col;
    }
}seg[maxn<<2];

///[s[u],e[u]]区间表示以u为根节点的子树的所有节点所在的区间
int s[maxn];
int e[maxn];
vector<int >vs;
void DFS(int u,int f)
{
    vs.push_back(u);
    s[u]=vs.size()-1;
    for(int i=head[u];~i;i=G[i].next)
    {
        int v=G[i].to;
        if(v != f)
            DFS(v,u);
    }
    e[u]=vs.size()-1;
}
void pushUp(int pos)
{
    seg[pos].sum=seg[ls(pos)].sum|seg[rs(pos)].sum;
}
void pushDown(int pos)
{
    int &lazy=seg[pos].lazy;
    if(lazy == -1)
        return ;

    seg[ls(pos)].Set(lazy);
    seg[rs(pos)].Set(lazy);

    lazy=-1;
}
void build(int l,int r,int pos)
{
    seg[pos]={l,r};
    seg[pos].lazy=-1;

    if(l == r)
    {
        seg[pos].sum=1ll<<a[vs[l]];
        return ;
    }

    int mid=seg[pos].mid();
    build(l,mid,ls(pos));
    build(mid+1,r,rs(pos));

    pushUp(pos);
}
void update(int pos,int l,int r,int col)
{
    if(seg[pos].l == l && seg[pos].r == r)
    {
        seg[pos].Set(col);
        return ;
    }
    pushDown(pos);

    int mid=seg[pos].mid();
    if(r <= mid)
        update(ls(pos),l,r,col);
    else if(l > mid)
        update(rs(pos),l,r,col);
    else
    {
        update(ls(pos),l,mid,col);
        update(rs(pos),mid+1,r,col);
    }
    pushUp(pos);
}
ll query(int pos,int l,int r)
{
    if(seg[pos].l == l && seg[pos].r == r)
        return seg[pos].sum;
    pushDown(pos);

    int mid=seg[pos].mid();
    if(r <= mid)
        return query(ls(pos),l,r);
    else if(l > mid)
        return query(rs(pos),l,r);
    else
        return query(ls(pos),l,mid)|query(rs(pos),mid+1,r);
}
void Solve()
{
    vs.clear();
    DFS(1,1);
    build(0,vs.size()-1,1);

    while(m--)
    {
        int op;
        scanf("%d",&op);
        if(op == 1)
        {
            int v,c;
            scanf("%d%d",&v,&c);
            update(1,s[v],e[v],c);///更新v节点对应的子树节点区间[s[v],e[v]]的颜色
        }
        else
        {
            int v;
            scanf("%d",&v);
            ll ans=query(1,s[v],e[v]);
            printf("%d
",popcount(ans));
        }
    }
}
void Init()
{
    num=0;
    mem(head,-1);
}
int main()
{
//    freopen("C:\Users\hyacinthLJP\Desktop\C++WorkSpace\in&&out\contest","r",stdin);
    Init();
    scanf("%d%d",&n,&m);
    for(int i=1;i <= n;++i)
        scanf("%d",a+i);
    for(int i=1;i < n;++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        addEdge(u,v);
        addEdge(v,u);
    }
    Solve();

    return 0;
}