线段树+Dfs序【CF620E】New Year Tree

Description

你有一棵以1为根的有根树，有n个点，每个节点初始有一个颜色c[i]。

有两种操作：

1 v c 将以v为根的子树中所有点颜色更改为c

2 v 查询以v为根的子树中的节点有多少种不同的颜色

Input

第一行,两个整数\(n,m\),分别代表有\(n\)个节点和\(m\)个操作。

第二行,共\(n\)个整数,代表每个节点的初始颜色\(c[i]\)

接下来\(n-1\)行,描述一条边。

接下来\(m\)行,代表每个操作。

Output

对于每个询问操作,输出一行。

刚开始以为是树剖?

结果发现只需要对每个子树操作。

线段树维护\(dfs\)序。

对于颜色呢？发现\(c[i] \leq 60\)

开$long  long $可以压成一个数。

因此我们将颜色压缩即可。

记得开$long  long $

虽然没出第二个样例,但我切了

代码

#include<cstdio>
#include<iostream>
#include<algorithm>
#define int long long 
#define R register

using namespace std;

const int gz=4e5+8;

inline void in(int &x)
{
    int f=1;x=0;char s=getchar();
    while(!isdigit(s)){if(s=='-')f=-1;s=getchar();}
    while(isdigit(s)){x=x*10+s-'0';s=getchar();}
    x*=f;
}

int head[gz],tot;

struct cod{int u,v;}edge[gz<<1];

inline void add(R int x,R int y)
{
    edge[++tot].u=head[x];
    edge[tot].v=y;
    head[x]=tot;
}

int dfn[gz],fdfn[gz],idx,size[gz];

void dfs(R int u,R int fa)
{
    dfn[u]=++idx;fdfn[idx]=u;size[u]=1;
    for(R int i=head[u];i;i=edge[i].u)
    {
        if(edge[i].v==fa)continue;
        dfs(edge[i].v,u);
        size[u]+=size[edge[i].v];
    }
}

int tr[gz<<2],c[gz],n,m;

bool tg[gz<<2];

#define ls o<<1
#define rs o<<1|1

inline void up(R int o)
{
    tr[o]=(tr[ls] | tr[rs]);
}

void build(R int o,R int l,R int r)
{
    if(l==r)
    {
        tr[o]=(1LL<<c[fdfn[l]]);
        return ;
    }
    R int mid=(l+r)>>1;
    build(ls,l,mid);
    build(rs,mid+1,r);
    up(o);
}

inline void down(R int o)
{
    if(tg[o])
    {
        tg[ls]=tg[rs]=tg[o];
        tr[ls]=tr[rs]=tr[o];
        tg[o]=false;
    }
}

void change(R int o,R int l,R int r,R int x,R int y,R int k)
{
    if(x<=l and y>=r){tr[o]=(1LL<<k);tg[o]=true;return;}
    down(o);
    R int mid=(l+r)>>1;
    if(x<=mid)change(ls,l,mid,x,y,k);
    if(y>mid)change(rs,mid+1,r,x,y,k);
    up(o);
}

int query(R int o,R int l,R int r,R int x,R int y)
{
    if(x<=l and y>=r)return tr[o];
    down(o);
    R int mid=(l+r)>>1;
    if(y<=mid)return query(ls,l,mid,x,y);
    else if(x>mid)return query(rs,mid+1,r,x,y);
    return (query(ls,l,mid,x,mid) | query(rs,mid+1,r,mid+1,y));
}

#define lowbit(o) o&-o 

inline int tquery(R int v)
{
    R int k=query(1,1,n,dfn[v],dfn[v]+size[v]-1);
    R int cnt=0;
    while(k) k-=lowbit(k),cnt++;
    return cnt;
}

signed main()
{
    in(n);in(m);
    for(R int i=1;i<=n;i++)in(c[i]);
    for(R int i=1,x,y;i<n;i++)
    {
        in(x),in(y);
        add(x,y);
        add(y,x);
    }
    dfs(1,0);build(1,1,n);
    for(R int i=1,opt,v,c;i<=m;i++)
    {
        in(opt);
        switch(opt)
        {
            case 1:in(v),in(c);change(1,1,n,dfn[v],dfn[v]+size[v]-1,c);break;
            case 2:in(v);printf("%lld\n",tquery(v));break;
        }
    }
}