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

 1 #include <iostream>
 2 #include <cstdio>
 3 #include <vector>
 4 #include <cstring>
 5 #include <algorithm>
 6 using namespace std;
 7 const int INF=0x3f3f3f3f;
 8 typedef long long ll;
 9 const int maxn=1e6+5;
10 
11 ll n,k,m;
12 ll ans;
13 ll a[maxn];
14 ll b[maxn];
15 ll c[maxn];
16 int degree[maxn];
17 
18 vector<int> G[maxn];
19 
20 int lowbit(int x)
21 {
22     return x&-x;
23 }
24 
25 int sum(int x)
26 {
27     int ret = 0;
28     while(x>0)
29     {
30         ret+=c[x]; x-=lowbit(x);
31     }
32     return ret;
33 }
34 
35 void add(int x, int d)
36 {
37     while(x<=m)
38     {
39         c[x]+=d;
40         x+=lowbit(x);
41     }
42 }
43 
44 void dfs(int u, int fa)
45 {
46     int update_pos=lower_bound(b+1,b+m+1,a[u])-b;
47     int query_pos;
48     if(a[u]) query_pos=lower_bound(b+1,b+m+1,k/a[u])-b;
49     else query_pos=m;
50     ans+=sum(query_pos);
51     add(update_pos,1);
52     for(int i=0;i<G[u].size();i++)
53     {
54         int v=G[u][i];
55         if(v==fa) continue;
56         dfs(v,u);
57     }
58     add(update_pos,-1);
59 }
60 
61 int main()
62 {
63     //freopen("in.txt","r",stdin);
64     int T;
65     scanf("%d",&T);
66     while(T--)
67     {
68         memset(c,0,sizeof(c));
69         memset(degree,0,sizeof(degree));
70         scanf("%lld%lld",&n,&k);
71         for(int i=1;i<=n;i++)  {G[i].clear();scanf("%lld",&a[i]);b[i]=a[i];}
72 
73         m=n;
74         for(int i=1;i<=n;i++)
75         {
76             if(a[i])  b[++m]=k/a[i];
77         }
78 
79         sort(b+1,b+m+1);
80         m=unique(b+1,b+m+1)-(b+1);
81 
82         for(int i=1;i<=n-1;i++)
83         {
84             int u,v;
85             scanf("%d%d",&u,&v);
86             G[u].push_back(v);
87             G[v].push_back(u);
88             degree[v]++;
89         }
90 
91         ans=0;
92         for(int i=1;i<=n;i++)
93             if(degree[i]==0)  dfs(i,-1);
94         printf("%lld\n",ans);
95     }
96     return 0;
97 }