題目連結:【HDU】5574 Colorful Tree
題目大意:對一個子樹染色,詢問一個子樹的顔色數。
題目分析: set 維護每種顔色所在的 dfs 序區間,修改均攤 nlogn 。
#include <bits/stdc++.h>
using namespace std ;
typedef long long LL ;
typedef pair < int , int > pii ;
#define clr( a , x ) memset ( a , x , sizeof a )
const int MAXN = ;
const int MAXE = ;
struct Edge {
int v , n ;
Edge () {}
Edge ( int v , int n ) : v ( v ) , n ( n ) {}
} ;
Edge E[MAXE] ;
int H[MAXN] , cntE ;
int dep[MAXN] ;
int pre[MAXN] ;
int top[MAXN] ;
int son[MAXN] ;
int siz[MAXN] ;
int tree_idx ;
int in[MAXN] ;
int ou[MAXN] ;
int dfs_idx ;
int n , q ;
set < pii > s , col[MAXN] ;
set < pii > :: iterator it , it1 ;
int c[MAXN] ;
int T[MAXN] ;
int setv[MAXN << ] ;
void addedge ( int u , int v ) {
E[cntE] = Edge ( v , H[u] ) ;
H[u] = cntE ++ ;
}
void dfs ( int u ) {
in[u] = ++ dfs_idx ;
siz[u] = ;
son[u] = ;
for ( int i = H[u] ; ~i ; i = E[i].n ) {
int v = E[i].v ;
if ( v == pre[u] ) continue ;
pre[v] = u ;
dep[v] = dep[u] + ;
dfs ( v ) ;
siz[u] += siz[v] ;
if ( siz[son[u]] < siz[v] ) son[u] = v ;
}
ou[u] = dfs_idx ;
}
void rebuild ( int u , int top_element ) {
top[u] = top_element ;
if ( son[u] ) rebuild ( son[u] , top_element ) ;
for ( int i = H[u] ; ~i ; i = E[i].n ) {
int v = E[i].v ;
if ( v != pre[u] && v != son[u] ) rebuild ( v , v ) ;
}
}
int lca ( int x , int y ) {
while ( top[x] != top[y] ) dep[top[x]] > dep[top[y]] ? x = pre[top[x]] : y = pre[top[y]] ;
return dep[x] < dep[y] ? x : y ;
}
void build ( int o , int l , int r ) {
setv[o] = ;
if ( l == r ) {
setv[o] = c[in[l]] ;
return ;
}
int m = l + r >> ;
build ( o << , l , m ) ;
build ( o << | , m + , r ) ;
}
void down ( int o ) {
if ( setv[o] ) {
setv[o << ] = setv[o << | ] = setv[o] ;
setv[o] = ;
}
}
void update ( int L , int R , int v , int o , int l , int r ) {
if ( L <= l && r <= R ) {
setv[o] = v ;
return ;
}
down ( o ) ;
int m = l + r >> ;
if ( L <= m ) update ( L , R , v , o << , l , m ) ;
if ( m < R ) update ( L , R , v , o << | , m + , r ) ;
}
int query ( int x , int o , int l , int r ) {
while ( l < r ) {
down ( o ) ;
int m = l + r >> ;
if ( x <= m ) o = o << , r = m ;
else o = o << | , l = m + ;
}
return setv[o] ;
}
void add ( int x , int v ) {
for ( int i = x ; i <= n ; i += i & -i ) T[i] += v ;
}
int sum ( int x , int ans = ) {
for ( int i = x ; i >= ; i -= i & -i ) ans += T[i] ;
return ans ;
}
int get ( int x ) {
int x1 = , x2 = ;
if ( it != col[c[x]].begin () ) {
it1 = it ;
-- it1 ;
x1 = lca ( it1->second , x ) ;
}
it1 = it ;
++ it1 ;
if ( it1 != col[c[x]].end () ) x2 = lca ( it1->second , x ) ;
if ( dep[x1] > dep[x2] ) return x1 ;
return x2 ;
}
void solve () {
int op , x , y ;
scanf ( "%d" , &n ) ;
cntE = dfs_idx = tree_idx = ;
s.clear () ;
for ( int i = ; i <= n ; ++ i ) {
H[i] = - ;
col[i].clear () ;
T[i] = ;
}
for ( int i = ; i < n ; ++ i ) {
scanf ( "%d%d" , &x , &y ) ;
addedge ( x , y ) ;
addedge ( y , x ) ;
}
dep[] = ;
dfs ( ) ;
rebuild ( , ) ;
for ( int i = ; i <= n ; ++ i ) {
scanf ( "%d" , &c[i] ) ;
col[c[i]].insert ( make_pair ( in[i] , i ) ) ;
s.insert ( make_pair ( in[i] , i ) ) ;
}
build ( , , n ) ;
for ( int i = ; i <= n ; ++ i ) {
int pre = ;
for ( it = col[i].begin () ; it != col[i].end () ; ++ it ) {
add ( it->first , ) ;
if ( pre ) add ( in[lca ( pre , it->second )] , - ) ;
pre = it->second ;
}
}
scanf ( "%d" , &q ) ;
for ( int i = ; i <= q ; ++ i ) {
scanf ( "%d%d" , &op , &x ) ;
pii tmp ( in[x] , x ) ;
if ( op == ) {
scanf ( "%d" , &y ) ;
while ( ) {
it = s.lower_bound ( tmp ) ;
if ( it == s.end () || it->first > ou[x] ) break ;
int u = it->second ;
it = col[c[u]].find ( *it ) ;
int x1 = get ( u ) ;
if ( x1 ) add ( in[x1] , ) ;
add ( in[u] , - ) ;
pii tmp1 = *it ;
s.erase ( tmp1 ) ;
col[c[u]].erase ( tmp1 ) ;
c[u] = y ;
}
c[x] = y ;
s.insert ( tmp ) ;
col[y].insert ( tmp ) ;
it = col[y].find ( tmp ) ;
int x1 = get ( x ) ;
if ( x1 ) add ( in[x1] , - ) ;
add ( in[x] , ) ;
update ( in[x] , ou[x] , y , , , n ) ;
} else {
int ans = sum ( ou[x] ) - sum ( in[x] - ) ;
it = s.lower_bound ( tmp ) ;
if ( it == s.end () ) ++ ans ;
else if ( it->second != x ) {
int color = query ( in[x] , , , n ) ;
it = col[color].lower_bound ( tmp ) ;
if ( it == col[color].end () || it->first > ou[x] ) ++ ans ;
}
printf ( "%d\n" , ans ) ;
}
}
}
int main () {
int T ;
scanf ( "%d" , &T ) ;
for ( int i = ; i <= T ; ++ i ) {
printf ( "Case #%d:\n" , i ) ;
solve () ;
}
return ;
}
壓縮後代碼:
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
typedef pair<int ,int >pii;
#define clr(a,x) memset(a,x,sizeof a)
const int MAXN=;
const int MAXE=;
struct Edge{
int v,n;
Edge(){}
Edge(int v,int n):v(v),n(n){}
}E[MAXE];
int H[MAXN],cntE;
int dep[MAXN],pre[MAXN],top[MAXN],son[MAXN],siz[MAXN],tree_idx;
int in[MAXN],ou[MAXN],dfs_idx;
set<pii>s,col[MAXN];
set<pii>::iterator it,it1;
int c[MAXN],T[MAXN],setv[MAXN<<],n,q;
void addedge(int u,int v){E[cntE]=Edge(v,H[u]);H[u]=cntE++;}
void dfs(int u){
in[u]=++dfs_idx;
siz[u]=;
son[u]=;
for(int i=H[u];~i;i=E[i].n){
int v=E[i].v;
if(v==pre[u])continue;
pre[v]=u;
dep[v]=dep[u]+;
dfs(v);
siz[u]+=siz[v];
if(siz[son[u]]<siz[v])son[u]=v;
}
ou[u]=dfs_idx;
}
void rebuild(int u,int top_element){
top[u]=top_element;
if(son[u])rebuild(son[u],top_element);
for(int i=H[u];~i;i=E[i].n)if(E[i].v!=pre[u]&&E[i].v!=son[u])rebuild(E[i].v,E[i].v);
}
int lca(int x,int y){
while(top[x]!=top[y])dep[top[x]]>dep[top[y]]?x=pre[top[x]]:y=pre[top[y]];
return dep[x]<dep[y]?x:y;
}
void build(int o,int l,int r){
setv[o]=;
if(l==r)setv[o]=c[in[l]];
else {
int m=l+r>>;
build(o<<,l,m);
build(o<<|,m+,r);
}
}
void down(int o){if(setv[o])setv[o<<]=setv[o<<|]=setv[o],setv[o]=;}
void update(int L,int R,int v,int o,int l,int r){
if(L<=l&&r<=R)setv[o]=v;
else {
down(o);
int m=l+r>>;
if(L<=m)update(L,R,v,o<<,l,m);
if(m<R)update(L,R,v,o<<|,m+,r);
}
}
int query(int x,int o,int l,int r){
for(int m=l+r>>;l<r;x<=m?(o=o<<,r=m):(o=o<<|,l=m+),m=l+r>>)down(o);
return setv[o];
}
void add(int x,int v){for(int i=x;i<=n;i+=i&-i)T[i]+=v;}
int sum(int x,int ans=){
for(int i=x;i>=;i-=i&-i)ans+=T[i];
return ans;
}
int get(int x,int x1=,int x2=){
if(it!=col[c[x]].begin())x1=lca((--(it1=it))->second,x);
if((++(it1=it))!=col[c[x]].end())x2=lca(it1->second,x);
if(dep[x1]>dep[x2])return x1;
return x2;
}
void solve(){
int op,x,y;
scanf("%d",&n);
cntE=dfs_idx=tree_idx=;
s.clear();
for(int i=;i<=n;++i)H[i]=-,col[i].clear(),T[i]=;
for(int i=;i<n;addedge(x,y),addedge(y,x),++i)scanf("%d%d",&x,&y);
dfs(dep[]=),rebuild(,);
for(int i=;i<=n;++i){
scanf("%d",&c[i]);
col[c[i]].insert(make_pair(in[i],i));
s.insert(make_pair(in[i],i));
}
build(,,n);
for(int i=,pre=;i<=n;++i,pre=){
for(it=col[i].begin();it!=col[i].end();++it){
add(it->first,);
if(pre)add(in[lca(pre,it->second)],-);
pre=it->second;
}
}
scanf("%d",&q);
for(int i=;i<=q;++i){
scanf("%d%d",&op,&x);
pii tmp(in[x],x);
if(op==){
scanf("%d",&y);
while(){
it=s.lower_bound(tmp);
if(it==s.end()||it->first>ou[x])break;
int u=it->second;
it=col[c[u]].find(*it);
int x1=get(u);
if(x1)add(in[x1],);
add(in[u],-);
pii tmp1=*it;
s.erase(tmp1);
col[c[u]].erase(tmp1);
c[u]=y;
}
c[x]=y;
s.insert(tmp);
col[y].insert(tmp);
it=col[y].find(tmp);
int x1=get(x);
if(x1)add(in[x1],-);
add(in[x],);
update(in[x],ou[x],y,,,n);
}else{
int ans=sum(ou[x])-sum(in[x]-);
it=s.lower_bound(tmp);
if(it==s.end())++ans;
else if(it->second!=x){
int color=query(in[x],,,n);
it=col[color].lower_bound(tmp);
if(it==col[color].end()||it->first>ou[x])++ans;
}
printf("%d\n",ans);
}
}
}
int main(){
int T;
scanf("%d",&T);
for(int i=;i<=T;++i){
printf("Case #%d:\n",i);
solve();
}
return ;
}