HDU 5877 Weak Pair dfs+树状数组

                                                 Weak Pair

Time Limit: 4000/2000 MS (Java/Others)    Memory Limit: 262144/262144 K (Java/Others)
Total Submission(s): 5280    Accepted Submission(s): 1528

 

Problem Description

You are given a rooted tree of N nodes, labeled from 1 to N. To the ith node a non-negative value ai is assigned.An ordered pair of nodes (u,v) is said to be weak if
  (1) u is an ancestor of v (Note: In this problem a node u is not considered an ancestor of itself);
  (2) au×av≤k.

Can you find the number of weak pairs in the tree?

 

 

Input

There are multiple cases in the data set.
  The first line of input contains an integer T denoting number of test cases.
  For each case, the first line contains two space-separated integers, N and k, respectively.
  The second line contains N space-separated integers, denoting a1 to aN.
  Each of the subsequent lines contains two space-separated integers defining an edge connecting nodes u and v , where node u is the parent of node v.

  Constrains:
  
  1≤N≤105
  
  0≤ai≤109
  
  0≤k≤1018

Output

For each test case, print a single integer on a single line denoting the number of weak pairs in the tree.

Sample Input

1 2 3 1 2 1 2

Sample Output

1

 

#include <bits/stdc++.h>

using namespace std;
typedef long long LL;
const int maxn = 1e5+100;
const int INF = 0x3f3f3f3f;

int n,T;
LL k;
LL a[maxn],b[maxn],c[maxn];
int d[maxn];
map<int,int> mp;
vector <int> G[maxn];
LL ans = 0;
int tot;
void dfs(int y)
{
  LL  v = k/a[y];
//  cout << a[y] << v << endl;
  int j = upper_bound(b,b+n,v)-b-1;
  for(int x = mp[b[j]]; x ; x -= x&-x)  ans += c[x];
  for(int x = mp[a[y]]; x <= tot; x += x&-x)   c[x]++;
  for(int i = 0; i < G[y].size(); i++)
     {
       dfs(G[y][i]);
       for(int x = mp[a[G[y][i]]]; x <= tot; x += x&-x) c[x]--;
     }
}
int main()
{
     scanf("%d",&T);
     while(T--)
   {
     mp.clear();
     scanf("%d%lld",&n,&k);
     for(int i = 0; i < n; i++)  {scanf("%lld",&a[i]); b[i] = a[i];G[i].clear();}
     tot = 1;
     sort(b,b+n);
     for(int i = 0; i < n; i++)
        if(mp.count(b[i]) == 0)  mp[b[i]] = tot++;

    memset(d,0,sizeof(d));
     for(int i = 0; i < n-1; i++)
       {
         int u,v;
         scanf("%d%d",&u,&v);
         G[u-1].push_back(v-1);
         d[v-1]++;
       }
     int root;
     for(int i = 0; i < n; i++)
       if(d[i] == 0) {root = i; break;}
    memset(c,0,sizeof(c));
    ans = 0LL;
  //  cout << root << endl;
    dfs(root);
    printf("%lld\n",ans);

 }
   return 0;
}