HDU - 4358 Boring counting (dsu on tree)

Boring counting: http://acm.hdu.edu.cn/showproblem.php?pid=4358

题意：

　　求一棵树上，每个节点的子节点中，同一颜色出现k次 的 个数。

思路：

　　由于是子树中出现了k次，sum+1。所以增加某种颜色的时候，如果这个颜色+1==k，那么sum++。如果删除的时候这个数+1 == k+1，那么sum--；

#include <algorithm>
#include  <iterator>
#include  <iostream>
#include   <cstring>
#include   <cstdlib>
#include   <iomanip>
#include    <bitset>
#include    <cctype>
#include    <cstdio>
#include    <string>
#include    <vector>
#include     <stack>
#include     <cmath>
#include     <queue>
#include      <list>
#include       <map>
#include       <set>
#include   <cassert>
#include <unordered_map>
using namespace std;
//#pragma GCC optimize(3)
//#pragma comment(linker, "/STACK:102400000,102400000")  //c++
// #pragma GCC diagnostic error "-std=c++11"
// #pragma comment(linker, "/stack:200000000")
// #pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,tune=native")
// #pragma GCC optimize("-fdelete-null-pointer-checks,inline-functions-called-once,-funsafe-loop-optimizations,-fexpensive-optimizations,-foptimize-sibling-calls,-ftree-switch-conversion,-finline-small-functions,inline-small-functions,-frerun-cse-after-loop,-fhoist-adjacent-loads,-findirect-inlining,-freorder-functions,no-stack-protector,-fpartial-inlining,-fsched-interblock,-fcse-follow-jumps,-fcse-skip-blocks,-falign-functions,-fstrict-overflow,-fstrict-aliasing,-fschedule-insns2,-ftree-tail-merge,inline-functions,-fschedule-insns,-freorder-blocks,-fwhole-program,-funroll-loops,-fthread-jumps,-fcrossjumping,-fcaller-saves,-fdevirtualize,-falign-labels,-falign-loops,-falign-jumps,unroll-loops,-fsched-spec,-ffast-math,Ofast,inline,-fgcse,-fgcse-lm,-fipa-sra,-ftree-pre,-ftree-vrp,-fpeephole2",3)

#define lson (l , mid , rt << 1)
#define rson (mid + 1 , r , rt << 1 | 1)
#define debug(x) cerr << #x << " = " << x << "\n";
#define pb push_back
#define pq priority_queue
#define max3(a,b,c) max(max(a,b),c)



typedef long long ll;
typedef unsigned long long ull;

typedef pair<ll ,ll > pll;
typedef pair<int ,int > pii;
typedef pair<int,pii> p3;

//priority_queue<int> q;//这是一个大根堆q
//priority_queue<int,vector<int>,greater<int> >q;//这是一个小根堆q
#define fi first
#define se second
//#define endl '\n'

#define OKC ios::sync_with_stdio(false);cin.tie(0)
#define FT(A,B,C) for(int A=B;A <= C;++A)  //用来压行
#define REP(i , j , k)  for(int i = j ; i <  k ; ++i)
//priority_queue<int ,vector<int>, greater<int> >que;

const ll mos = 0x7FFFFFFF;  //2147483647
const ll nmos = 0x80000000;  //-2147483648
const int inf = 0x3f3f3f3f;       
const ll inff = 0x3f3f3f3f3f3f3f3f; //18
const int mod = 1e9+7;
const double esp = 1e-8;
const double PI=acos(-1.0);



template<typename T>
inline T read(T&x){
    x=0;int f=0;char ch=getchar();
    while (ch<'0'||ch>'9') f|=(ch=='-'),ch=getchar();
    while (ch>='0'&&ch<='9') x=x*10+ch-'0',ch=getchar();
    return x=f?-x:x;
}


/*-----------------------showtime----------------------*/
                const int maxn = 1e5+9;
                vector<int>g[maxn];
                // vector<int>v[maxn];
                int col[maxn];
                unordered_map<ll,int>cnt;
                ll ans[maxn];
                bool big[maxn];
                int n,k;
                int sz[maxn];

                void dfs1(int v,int fa){
                    sz[v] = 1;
                    for(int i=0; i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa){
                            dfs1(u,v);
                            sz[v]+=sz[u];
                        }
                    }
                }
                ll sum = 0;
                // priority_queue<int>que;
                void add(int v, int fa,int x){
                    cnt[col[v]] += x;
                    // if(v==1)cout<<"$$"<<cnt[col[v]]<<endl;
                    if(cnt[col[v]] == k) sum++;
                    else if(cnt[col[v]] == k+1) sum--;

                    for(int i=0;i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa && big[u]==0){
                            add(u,v,x);
                        }
                    }
                }

                void dfs(int v,int fa,int keep){
                    int mx = -1,bigChild = -1;

                    for(int i=0; i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa && sz[u] > mx)
                            mx = sz[u], bigChild = u;
                    }

                    for(int i=0; i<g[v].size(); i++){
                        int u = g[v][i];
                        if(u!=fa && u!=bigChild){
                            dfs(u,v,0);
                        }
                    }

                    if(bigChild!=-1){
                        dfs(bigChild,v,1),big[bigChild] = 1;
                    }

                    add(v,fa,1);         
                    ans[v] = sum;

                    if(bigChild!=-1) big[bigChild] = 0;
                    if(keep==0){
                        add(v,fa,-1);
                        sum = 0;
                    }

                }
                void solve(){
                    sum = 0;
                    read(n),read(k);
                    for(int i=1; i<=n; i++)g[i].clear(),big[i] = 0;
                        cnt.clear();
                    for(int i=1; i<=n; i++)read(col[i]);
                        for(int i=1; i<n; i++){
                            int u,v;
                            read(u),read(v);
                            g[u].pb(v);
                            g[v].pb(u);
                        }
                        dfs1(1,0);
                        dfs(1,0,1);
                        int q;read(q);
                        while(q--){
                            int x;read(x);
                            printf("%lld\n", ans[x]);
                        }
                }
int main(){
                int cas;scanf("%d", &cas);
                for(int i=1; i<=cas;i++){
                    printf("Case #%d:\n",i);
                    solve();
                    if(i<cas)puts("");
                }
                return 0;
}