2015
05-23

Boring counting

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?

The first line of the input contains an integer 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 )
Then come N integers in the second line, they are the weights of vertice 1 to N. ( 0 <= weight <= 109 )
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)
For next Q lines, each with an integer u, as the root of the subtree described above.

The first line of the input contains an integer 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 )
Then come N integers in the second line, they are the weights of vertice 1 to N. ( 0 <= weight <= 109 )
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)
For next Q lines, each with an integer u, as the root of the subtree described above.

1
3 1
1 2 2
1 2
1 3
3
2
1
3

Case #1:
1
1
1

Round #136 (Div. 2) D. Little Elephant and Array

#include <iostream>
#include <cstdio>
#include <cstring>
#include <vector>
#include <map>
#include <algorithm>
using namespace std;
#define LL(x) (x<<1)
#define RR(x) (x<<1|1)
const int N=1e5+5;
struct Query
{
int st,ed,id;
Query(){}
Query(int a,int b,int c){st=a;ed=b;id=c;}
bool operator < (const Query &b) const
{
return ed<b.ed;
}
};
struct node
{
int lft,rht,sum;
int mid(){return lft+(rht-lft)/2;}
};
struct Segtree
{
node tree[N*4];
void build(int lft,int rht,int ind)
{
tree[ind].lft=lft;  tree[ind].rht=rht;
tree[ind].sum=0;
if(lft!=rht)
{
int mid=tree[ind].mid();
build(lft,mid,LL(ind));
build(mid+1,rht,RR(ind));
}
}
void updata(int pos,int ind,int valu)
{
tree[ind].sum+=valu;
if(tree[ind].lft==tree[ind].rht) return;
else
{
int mid=tree[ind].mid();
if(pos<=mid) updata(pos,LL(ind),valu);
else updata(pos,RR(ind),valu);
}
}
int query(int be,int end,int ind)
{
int lft=tree[ind].lft,rht=tree[ind].rht;
if(be<=lft&&rht<=end) return tree[ind].sum;
else
{
int mid=tree[ind].mid();
int sum1=0,sum2=0;
if(be<=mid) sum1=query(be,end,LL(ind));
if(end>mid) sum2=query(be,end,RR(ind));
return sum1+sum2;
}
}
}seg;
int data1[N],data2[N];
int low[N],high[N],res[N],idx;
vector<Query> q;
map<int,int> imap;

void dfs(int u,int fa)
{
low[u]=++idx;
data2[idx]=data1[u];
{
if(v==fa) continue;
dfs(v,u);
}
high[u]=idx;
}
void init(int n)
{
idx=0;
imap.clear(); q.clear();
for(int i=0;i<=n;i++)
{
}
}
int main()
{
int t,t_cnt=0;
scanf("%d",&t);
while(t--)
{
int n,k,a,b,m,sc=0;
scanf("%d%d",&n,&k);
init(n);
for(int i=1;i<=n;i++)
{
scanf("%d",&data1[i]);
if(imap.find(data1[i])==imap.end())
imap.insert(make_pair(data1[i],sc++));
}
for(int i=1;i<n;i++)
{
scanf("%d%d",&a,&b);
}
dfs(1,-1);
scanf("%d",&m);
for(int i=0;i<m;i++)
{
scanf("%d",&a);
q.push_back(Query(low[a],high[a],i));
}
sort(q.begin(),q.end());
int ind=0;
seg.build(1,idx,1);
for(int i=1;i<=n;i++)
{
int id=imap[data2[i]];
pos[id].push_back(i);
int cnt=(int)pos[id].size();
if(cnt>=k)
{
if(cnt>k)
seg.updata(pos[id][cnt-k-1],1,-2);
if(cnt>k+1)
seg.updata(pos[id][cnt-k-2],1,1);
seg.updata(pos[id][cnt-k],1,1);
}
while(ind<m&&q[ind].ed==i)
{
res[q[ind].id]=seg.query(q[ind].st,q[ind].ed,1);
ind++;
}

}
if(t_cnt!=0) puts("");
printf("Case #%d:\n",++t_cnt);
for(int i=0;i<m;i++) printf("%d\n",res[i]);
}
return 0;
}