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

时间:2023-03-09 20:08:18
牛客练习赛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';
}
}
}