HDU-5877-Weak Pair（离散+树状数组）

链接：http://acm.hdu.edu.cn/showproblem.php?pid=5877

题意：

给定一棵树求解满足以下条件的点对个数

1，对于(u,v)，u为v的祖先节点

2，对于(u,v)，有au*av<=k

题解：

DFS过程中维护a[i]的bit，查找k/a[i]，数据只有10W离散化以下就好。

CODE：

#include <bits/stdc++.h>
//#pragma comment(linker, "/STACK:1024000000,1024000000")
using namespace std;
#define INF 0x3f3f3f3f
#define LL long long
#define bug cout<<"bug"<<endl
const int MAXN = 200007;
const int MOD = 1e9 + 9;
long long a[MAXN],b[MAXN],k;
int num[MAXN],sum[MAXN],head[MAXN],ind,cnt;
int deep[MAXN];
struct node
{
int v,next;
}edge[MAXN];
void add_edge(int u, int v)
{
edge[ind].v=v;
edge[ind].next=head[u];
head[u]=ind++;
}
void add(int i, int x)
{
while(i<=MAXN)
{
sum[i]+=x;
i+=i&-i;
}
}
int get_sum(int i)
{
int ans=0;
while(i>0)
{
ans+=sum[i];
i-=i&-i;
}
return ans;
}
long long all;
void DFS(int u)
{
int l=lower_bound(b,b+cnt,k/a[u])-b+1;
all+=get_sum(l);
int r=lower_bound(b,b+cnt,a[u])-b+1;
add(r,1);
for(int i=head[u]; i+1; i=edge[i].next)
DFS(edge[i].v);
add(r,-1);
}
int main()
{
int T,n;
scanf("%d",&T);
while(T--)
{
cnt=all=ind=0;
memset(head,-1,sizeof(head));
memset(deep,0,sizeof(deep));
scanf("%d",&n);
scanf("%I64d",&k);
for(int i=1; i<=n; ++i)
{
scanf("%I64d",&a[i]);
b[cnt++]=a[i];
b[cnt++]=k/a[i];
}
int u,v;
sort(b,b+cnt);
int temp=unique(b,b+cnt)-b;
for(int i=1; i<n; ++i)
{
scanf("%d%d",&u,&v);
add_edge(u,v);
deep[v]++;
}
for(int i=1; i<=n; ++i)
if(deep[i]==0)
DFS(i);
cout<<all<<endl;
}
return 0;
}
/*
2
2 3
1 2
1 2

3 3
1 2 3
1 2
2 3

*/