牛客练习赛47 E DongDong数颜色 （树状数组维护区间元素种类数）

链接：https://ac.nowcoder.com/acm/contest/904/E

来源：牛客网

DongDong数颜色

时间限制：C/C++ 1秒，其他语言2秒

空间限制：C/C++ 524288K，其他语言1048576K

64bit IO Format: %lld

题目描述

DongDong是个喜欢数颜色的女孩子，她已经熟练地掌握了在序列上数颜色的操作，现在她开始学习如何在树上数颜色，现在给定一个n个点,n-1条边的树形图（视1号店为根），每个点有一个颜色，每次询问以x为根的子树中有多少种不同的颜色，DongDong轻松地解决了这个问题，但她想考考会编程的你。

输入描述:

第一行两个整数n,m

第二行n个整数，表示每个点的颜色

接下来n-1行每行u,v,表示存在一条从u到v的双向边（保证最终图形是树形图）

2<=n<=100000,1<=m,color<=n，

输出描述:

共m行：每行输出相应询问的答案

示例1

输入

复制

4 3

1 1 2 3

1 2

2 3

1 4

1

2

4

输出

复制

3

2

1

思路：

以一个节点为根的子树的节点在dfs序中是一个连续的区间，我们可以用树状数组维护dfs序列，并对每一个节点为根的子树区间进行以右端点升序的方式排序，

排序后我们扫描区间，用vis[i] 维护i这个元素的上一次出现的位置，如果没出现则为0，出现了则记录位置。当再一次出现该数字时，把上一个位置的贡献给删除，当前位置加上贡献。

为什么这可以这样写？因为我们已经对区间右端点排序。

还有dsu on tree 的写法，可以看我这个博客： https://www.cnblogs.com/qieqiemin/p/11310304.html

细节见代码：

#include <iostream>

#include <cstdio>

#include <cstring>

#include <algorithm>

#include <cmath>

#include <queue>

#include <stack>

#include <map>

#include <set>

#include <vector>

#include <iomanip>

#define ALL(x) (x).begin(), (x).end()

#define sz(a) int(a.size())

#define all(a) a.begin(), a.end()

#define rep(i,x,n) for(int i=x;i<n;i++)

#define repd(i,x,n) for(int i=x;i<=n;i++)

#define pii pair<int,int>

#define pll pair<long long ,long long>

#define gbtb ios::sync_with_stdio(false),cin.tie(0),cout.tie(0)

#define MS0(X) memset((X), 0, sizeof((X)))

#define MSC0(X) memset((X), '\0', sizeof((X)))

#define pb push_back

#define mp make_pair

#define fi first

#define se second

#define eps 1e-6

#define gg(x) getInt(&x)

#define chu(x) cout<<"["<<#x<<" "<<(x)<<"]"<<endl

using namespace std;

typedef long long ll;

ll gcd(ll a, ll b) {return b ? gcd(b, a % b) : a;}

ll lcm(ll a, ll b) {return a / gcd(a, b) * b;}

ll powmod(ll a, ll b, ll MOD) {ll ans = 1; while (b) {if (b % 2)ans = ans * a % MOD; a = a * a % MOD; b /= 2;} return ans;}

inline void getInt(int* p);

const int maxn = 100010;

const int inf = 0x3f3f3f3f;

/*** TEMPLATE CODE * * STARTS HERE ***/

int q;

int n;

std::vector<int> son[maxn];

int in[maxn];

int out[maxn];

int tree[maxn];

int lowbit(int x)

{

    return -x&x;

}

void add(int x,int val)

{

    while(x<maxn)

    {

        tree[x]+=val;

        x+=lowbit(x);

    }

}

int ask(int x)

{

    int res=0;

    while(x)

    {

        res+=tree[x];

        x-=lowbit(x);

    }

    return res;

}

int tot=0;

int num[maxn];

void dfs(int x,int pre)

{

    in[x]=++tot;

    num[tot]=x;

    for(auto y:son[x])

    {

        if(y!=pre)

        {

            dfs(y,x);

        }

    }

    out[x]=tot;

}

int a[maxn];

int vis[maxn];

struct node

{

    int l,r;

    int x;

    int id;

}b[maxn];

bool cmp(node aa,node bb)

{

    return aa.r<bb.r;

}

int ans[maxn];

int main()

{

    //freopen("D:\\code\\text\\input.txt","r",stdin);

    //freopen("D:\\code\\text\\output.txt","w",stdout);

    gbtb;

    cin>>n>>q;

    repd(i,1,n)

    {

        cin>>a[i];

    }

    int u,v;

    repd(i,2,n)

    {

        cin>>u>>v;

        son[u].pb(v);

        son[v].pb(u);

    }

    dfs(1,-1);

    repd(i,1,n)

    {

        b[i].l=in[i];

        b[i].r=out[i];

        b[i].x=a[i];

        b[i].id=i;

    }

    sort(b+1,b+1+n,cmp);

    int p=1;

    repd(i,1,n)

    {

        while(p<=b[i].r)

        {

            int y=num[p];

            if(vis[a[y]])

            {

                add(vis[a[y]],-1);

                add(p,1);

                vis[a[y]]=p;

            }else

            {

                vis[a[y]]=p;

                add(p,1);

            }

            p++;

        }

        ans[b[i].id]=ask(b[i].r)-ask(b[i].l-1);

    }

    while(q--)

    {

        int x;

        cin>>x;

        cout<<ans[x]<<endl;

    }

    return 0;

}

inline void getInt(int* p) {

    char ch;

    do {

        ch = getchar();

    } while (ch == ' ' || ch == '\n');

    if (ch == '-') {

        *p = -(getchar() - '0');

        while ((ch = getchar()) >= '0' && ch <= '9') {

            *p = *p * 10 - ch + '0';

        }

    }

    else {

        *p = ch - '0';

        while ((ch = getchar()) >= '0' && ch <= '9') {

            *p = *p * 10 + ch - '0';

        }

    }

}