Boring counting
Time Limit: 6000/3000 MS (Java/Others)
Memory Limit: 98304/98304 K (Java/Others)
Problem Description
In this problem we consider a rooted tree with N vertices. The vertices are numbered from 1 to N, and vertex 1 represents the root. There are integer weights on each vectice. Your task is to answer a list of queries, for each query, please tell us among all the vertices in the subtree rooted at vertice u, how many different kinds of weights appear exactly K times?
Input
The first line of the input contains an integer
T(T<=5)
T
(
T
<=
5
)
, indicating the number of test cases.
For each test case, the first line contains two integers N and K, as described above.
(1<=N<=105,1<=K<=N)
(
1
<=
N
<=
10
5
,
1
<=
K
<=
N
)
Then come N integers in the second line, they are the weights of vertice 1 to N.
(0<=weight<=109)
(
0
<=
w
e
i
g
h
t
<=
10
9
)
For next N-1 lines, each line contains two vertices u and v, which is connected in the tree.
Next line is a integer Q, representing the number of queries.
(1<=Q<=105)
(
1
<=
Q
<=
10
5
)
For next Q lines, each with an integer u, as the root of the subtree described above.
Output
For each test case, output “Case #X:” first, X is the test number. Then output Q lines, each with a number – the answer to each query.
Seperate each test case with an empty line.
Sample Input
1
3 1
1 2 2
1 2
1 3
3
2
1
3
Sample Output
Case #1:
1
1
1
思路:树上启发式合并貌似很好做?
#include<bits/stdc++.h>
using namespace std;
const int MAX=1e5+10;
vector<int>e[MAX];
vector<int>v;
int a[MAX],siz[MAX],L[MAX],R[MAX],son[MAX],num[MAX],all;
void dfs(int k,int fa)
{
siz[k]=1;
L[k]=++all;
num[all]=a[k];
son[k]=0;
for(int i=0;i<e[k].size();i++)
{
int nex=e[k][i];
if(nex==fa)continue;
dfs(nex,k);
siz[k]+=siz[nex];
if(siz[son[k]]<siz[nex])son[k]=nex;
}
R[k]=all;
}
int ans[MAX];
int cnt[MAX];
int tot;
int n,m;
void cal(int k,int fa,int tp)//tp=0/1表示k不是/是重儿子
{
for(int i=0;i<e[k].size();i++)
{
int nex=e[k][i];
if(nex==fa||nex==son[k])continue;//优先遍历轻儿子
cal(nex,k,0);
}
if(son[k])cal(son[k],k,1);//最后遍历重儿子,此时重儿子的信息已被存储下来
cnt[a[k]]++;
if(cnt[a[k]]==m)tot++;
if(cnt[a[k]]==m+1)tot--;
for(int i=0;i<e[k].size();i++)//暴力统计轻儿子对答案的影响
{
int nex=e[k][i];
if(nex==fa||nex==son[k])continue;
for(int j=L[nex];j<=R[nex];j++)
{
cnt[num[j]]++;
if(cnt[num[j]]==m)tot++;
if(cnt[num[j]]==m+1)tot--;
}
}
ans[k]=tot;
if(tp==0)//如果当前节点不是重儿子,清除信息
{
tot=0;
for(int j=L[k];j<=R[k];j++)cnt[num[j]]--;
}
}
int main()
{
int T;
cin>>T;
for(int cas=1;cas<=T;cas++)
{
scanf("%d%d",&n,&m);
v.clear();
for(int i=1;i<=n;i++)
{
scanf("%d",&a[i]);
v.push_back(a[i]);
}
sort(v.begin(),v.end());
v.erase(unique(v.begin(),v.end()),v.end());//将点的权值离散化
for(int i=1;i<=n;i++)a[i]=lower_bound(v.begin(),v.end(),a[i])-v.begin()+1;
for(int i=1;i<=n;i++)e[i].clear();
for(int i=1;i<=n;i++)cnt[i]=0;
for(int i=1;i<n;i++)
{
int x,y;
scanf("%d%d",&x,&y);
e[x].push_back(y);
e[y].push_back(x);
}
all=0;
dfs(1,0);
tot=0;
cal(1,0,1);
int Q;
scanf("%d",&Q);
printf("Case #%d:\n",cas);
while(Q--)
{
int x;
scanf("%d",&x);
printf("%d\n",ans[x]);
}
if(cas<T)printf("\n");
}
return 0;
}