codeforces dfs序+线段树+bitset

题目链接

http://codeforces.com/problemset/problem/620/E

题意

给定一个树，现在有60种颜色，要做如下操作
1 v c：将v及其子树然成c颜色
2 v：查询v及其子树中含有多少不同颜色

思路

用dfs序存储子树，bitset代表子树中含有的颜色。

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#include<vector>
#include<string>
#include<queue>
#include<bitset>
#include<stack>
#include<set>
#include<map>
#define ll long long
using namespace std;
const int INF = ( 2e9 ) + 2;
const ll maxn = 4e5+100;
vector<int> g[maxn];
int col[maxn];
int p1[maxn],p2[maxn],con[maxn];
int index;
bitset<61> z,o;
struct node
{
    int l,r,mark;
    bitset<61> p;
}t[maxn*4];
void build(int root,int l,int r)
{
    t[root].l=l;
    t[root].r=r;
    t[root].mark=0;
    if(l==r)
    {
        t[root].p[ col[con[l]] ]=1; // dfs序为l的结点的颜色 
        return;
    }
    int mid=(l+r)>>1;
    build(root*2,l,mid);
    build(root*2+1,mid+1,r);
    t[root].p=t[root*2].p | t[root*2+1].p;
}
void pushdown(int root)
{
    if(t[root].mark==1)
    {
        t[root*2].p.reset();
        t[root*2].p = t[root*2+1].p =t[root].p;
        t[root*2].mark= t[root*2+1].mark = 1;
        t[root*2+1].p.reset();  
        t[root].mark=0;
    }
}
void update(int root,int ul,int ur,int c)
{
    int l=t[root].l,r=t[root].r;
    if(ul>r||ur<l)return;
    if(ul<=l&&ur>=r)
    {
        t[root].p.reset();
        t[root].p[c]=1;
        t[root].mark=1;
        return;
    }
    int mid = (l+r)>>1;
    pushdown(root);
    if(ur<=mid) 
    update(root*2,ul,ur,c);
    else if(ul>mid)
    update(root*2+1,ul,ur,c);
    else
    {
        update(root*2,ul,ur,c);
        update(root*2+1,ul,ur,c);
    }
    t[root].p = t[root*2].p | t[root*2+1].p;
}
bitset<61> query(int root,int ql,int qr)
{
    int l=t[root].l,r=t[root].r;
    if(ql>r||qr<l)return o;
    if(ql<=l&&qr>=r)
    {
        return t[root].p;
    }
    int mid = (l+r)>>1;
    pushdown(root);
    if(qr<=mid)
    return query(root*2,ql,qr);
    else if(ql>mid)
    return query(root*2+1,ql,qr);
    else
    {
        return query(root*2,ql,qr) | query(root*2+1,ql,qr);
    }
}
void dfs(int u,int fa)
{
    p1[u]=++index;
    con[index]=u;
    for(int i=0,L=g[u].size();i<L;i++)
    {
        int v=g[u][i];
        if(v==fa)continue;
        dfs(v,u);
    }
    p2[u]=index;
}
int main()
{
    int n,m,u,v;
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)
        scanf("%d",&col[i]);
    for(int i=1;i<=n-1;i++)
    {
        scanf("%d%d",&u,&v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,-1);
    build(1,1,n);
    for(int i=1;i<=60;i++)o[i]=1;
    int op;
    for(int i=1;i<=m;i++)
    {
        scanf("%d",&op);
        if(op==1)
        {
            scanf("%d%d",&v,&u);
            update(1,p1[v],p2[v],u);
        }
        else
        {
            scanf("%d",&v);
            int cnt=0;
            bitset<61> get=query(1,p1[v],p2[v]);
            printf("%d\n",get.count());
        }
    }
}