Codeforces 375D Tree and Queries

Codeforces 375D Tree and Queries

數據結構，dsu on tree

題意

給一棵樹，每個節點有顏色，給一些查詢（k，v），問k及其子樹中，出現大於等於v次的顏色有多少種。

思路

dsu on tree。複雜度O(nlog2n) ，dsu的nlogn加上樹狀數組的logn。

dsu的過程中用樹狀數組維護出現t次的顏色的種數。dfs完一個節點處理他的所有詢問。

代碼

#include<bits\stdc++.h>
#define M(a,b) memset(a,b,sizeof(a))
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
using namespace std;
const int MAXN=100007;
const int
typedef long long LL;
struct Edge
{
    int to, ne;
}e[MAXN<<1];
int head[MAXN], edgenum;
int val[MAXN];
struct Query
{
    int k, ne;int id;
}qn[MAXN];
int qhead[MAXN], qam;
int sz[MAXN], hson[MAXN];
int C[MAXN];
inline int lowbit(int i) { return i&(-i); }
void change(int pos, int val)
{
    for(int i=pos;i<MAXN;i+=lowbit(i))
        C[i]+=val;
}
int query(int pos)
{
    int ans=0;
    for(int i=pos;i>=1;i-=lowbit(i))
        ans+=C[i];
    return ans;
}
void findson(int u, int fa)
{
    sz[u]=1;
    for(int i=head[u];~i;i=e[i].ne)
    {
        int to=e[i].to;
        if(to!=fa)
        {
            findson(to, u);
            if(!hson) hson[u]=to;
            else if(sz[hson[u]]<sz[to]) hson[u]=to;
            sz[u]+=sz[to];
        }
    }
}
int cnt[MAXN], hs;
void suan(int u, int fa, int v)
{
    if(cnt[val[u]]) change(cnt[val[u]], -1);
    cnt[val[u]]+=v;
    if(cnt[val[u]]) change(cnt[val[u]], 1);
    for(int i=head[u];~i;i=e[i].ne)
        if(e[i].to!=fa&&e[i].to!=hs)
            suan(e[i].to, u, v);
}
int res[MAXN];
void dfs(int u, int fa, int kp)
{
    for(int i=head[u];~i;i=e[i].ne)
    {
        int to=e[i].to;
        if(to!=fa&&to!=hson[u])
            dfs(to, u, 0);
    }
    if(hson[u]) dfs(hson[u], u, 1), hs=hson[u];
    suan(u, fa, 1);hs=0;
    for(int i=qhead[u];~i;i=qn[i].ne)
    {
        int x=query(MAXN-1);
        int y=query(qn[i].k-1);
        res[qn[i].id]=x-y;
    }
    if(!kp) suan(u, fa, -1);
}
int main()
{
    int n;
    while(scanf("%d", &n)==1)
    {
        int q;scanf("%d", &q);
        for(int i=1;i<=n;i++) scanf("%d", &val[i]);
        M(head, -1), M(qhead, -1);edgenum=qam=0;
        for(int i=1;i<n;i++)
        {
            int p, q;scanf("%d%d", &p, &q);
            e[edgenum].to=q, e[edgenum].ne=head[p], head[p]=edgenum++;
            e[edgenum].to=p, e[edgenum].ne=head[q], head[q]=edgenum++;
        }
        for(int i=1;i<=q;i++)
        {
            int t1, t2;scanf("%d%d", &t1, &t2);
            qn[qam].k=t2, qn[qam].ne=qhead[t1], qn[qam].id=i, qhead[t1]=qam++;
        }
        findson(1, -1);
        dfs(1, 0, 0);
        for(int i=1;i<=q;i++)
            printf("%d\n", res[i]);
    }
    //system("pause");
    return 0;
}