HDU5877 Weak Pair(樹狀數組+dfs)

Weak Pair

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

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 nodesu 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

 

Source

2016 ACM/ICPC Asia Regional Dalian Online

題意：有一棵樹，N個結點，N-1條邊，每個結點有權值ai，如果存在一個點對，由結點和它的祖先組成，
並且它們相乘的結果<=K,那麼這個點對就稱爲weak，求有多少個這樣的點對。
思路：先將ai和K/ai離散化，然後找到根節點遍歷樹，每插入一個結點前，用樹狀數組求當前小於K/ai的結點
有多少個並加入結果，然後再將樹狀數組ai的位置+1，每次回溯前將樹狀數組ai的位置-1。注意ai可能爲0，所
以K/ai要處理一下。

#include <bits/stdc++.h>
using namespace std;
#define ll long long int
#define maxn 100010
ll bit[maxn*2], a[maxn], b[maxn*2], ans, K;
int N, n, tot, head[maxn];
bool vis[maxn];
void add(ll bit[], int i, ll x){
    while(i <= n){
        bit[i] += x;
        i += i&-i;
    }
}
ll sum(ll bit[], int i){
    ll s = 0;
    while(i > 0){
        s += bit[i];
        i -= i&-i;
    }
    return s;
}
struct Edge{
    int to, next;
}edge[maxn];
void addedge(int u, int v){
    edge[tot].to = v;
    edge[tot].next = head[u];
    head[u] = tot++;
}
void dfs(int now){
    vis[now] = 1;
    //printf("%d ", now);
    ll loc1, loc2;
    if(a[now] == 0) loc1 = n;
    else {
        ll point = K/a[now];
        //printf("%I64d ", point);
        loc1 = lower_bound(b, b+n, point)-b+1;
    }
    loc2 = lower_bound(b, b+n, a[now])-b+1;
    //printf("%I64d\n", sum(bit, loc1));
    ans += sum(bit, loc1);
    add(bit, loc2, 1);
    for(int i = head[now];i != -1;i = edge[i].next){
        int son = edge[i].to;
        if(vis[son]) continue;
        dfs(son);
    }
    add(bit, loc2, -1);
}
int main()
{
    int i, T, top;
    scanf("%d", &T);
    while(T--){
        scanf("%d %I64d", &N, &K);
        top = 0;
        for(i = 1;i <= N;i++){
            scanf("%I64d", &a[i]);
            b[top++] = a[i];
            if(a[i] != 0) b[top++] = K/a[i];
        }
        sort(b, b+top);
        n = unique(b, b+top)-b;
        int u, v;
        memset(head, -1, sizeof head);
        memset(bit, 0, sizeof bit);
        memset(vis, 0, sizeof vis);
        tot = 0;
        ans = 0;
        for(i = 0;i < N-1;i++){
            scanf("%d %d", &u, &v);
            addedge(u, v);
            vis[v] = 1;
        }
        int start;
        for(i = 1;i <= N;i++)
            if(!vis[i]){
                start = i;
                break;
            }
        memset(vis, 0, sizeof vis);
        dfs(start);
        printf("%I64d\n", ans);
    }
}