HDU 5877 Weak Pair 【dfs】【树状数组】【离散化】

题目：点击打开链接

题意：给出一棵有根树，找出多少对点，满足两点是祖先关系且权值之积小于等于k。

思路：可以用链式前向星保存边，建立一棵树，容易想到au*av<=k等价为au<=k/av，所以我们可以先把各权值被k整除，保存进数组中，对于每一个子节点，其祖先结点都肯定已经遍历过，所以可以用树状数组实现，由于是树形结构，所以必定也要用dfs遍历整棵树，每次遍历到一个结点，先把其权值离散化，然后保存到树状数组中，再遍历其子节点，利用树状数组找到所有比子节点的权值被k整除的数的个数，累加到答案中，遍历完所有子节点后将其从树状数组中删除，答案就是所有的个数。

代码：

#include <stdio.h>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <queue>
#include <set>
#include <vector>
#include <map>
#define PR pair<int,int>
#define MP make_pair
#define fi first
#define se second
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define sqr(x) ((x)*(x))
#define ll long long
#define ull unsigned __int64
const ll INF = 1e18+10;
const int inf=0x3f3f3f3f;
const int M=1000010;
const int N=100100;
const int MOD=10007;
const double eps=1e-10;
const double pi=acos(-1.0);
using namespace std;
struct node
{
int v,nxt;
}edge
;
int head
,cnt,n,u,v,vis
;
ll a
,k,ans,s[N*2];
vector<ll> vt;
int getid(ll x) {return lower_bound(vt.begin(),vt.end(),x)-vt.begin()+1;}
void addedge(int u,int v)
{
edge[cnt].v=v,edge[cnt].nxt=head[u],head[u]=cnt++;
}
void add(int p,int v)
{
for(int i=p;i<=n+1;i+=i&(-i))
s[i]+=v;
}
ll query(int p)
{
ll ans=0;
for(int i=p;i>0;i-=i&(-i))
ans+=s[i];
return ans;
}
void dfs(int u)
{
add(getid(a[u]),1);
for(int i=head[u];i!=-1;i=edge[i].nxt)
{
int v=edge[i].v;
if(a[v]==0) ans+=query(getid(INF));
else ans+=query(getid(k/a[v]));
dfs(v);
}
add(getid(a[u]),-1);
}
int main()
{
int casT,i;
scanf("%d",&casT);
while(casT--)
{
vt.clear(),ans=cnt=0;
memset(vis,0,sizeof(vis));
memset(head,-1,sizeof(head));
scanf("%d%I64d",&n,&k);
for(i=1;i<=n;i++)
{
scanf("%I64d",&a[i]);
if(a[i]==0) vt.push_back(INF);
else vt.push_back(k/a[i]);
}
sort(vt.begin(),vt.end()),vt.erase(unique(vt.begin(),vt.end()),vt.end());
for(i=1;i<n;i++)
{
scanf("%d%d",&u,&v);
addedge(u,v);
vis[v]=1;
}
for(i=1;i<=n;i++) if(vis[i]==0) dfs(i);
printf("%I64d\n",ans);
}
}