HDU 5877 - Weak Pair

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

题意：给出一棵 n 个节点的树，节点的点权为 a[i]，要求找出有多少组边满足一下条件：
1：u是v的祖先且u != v
2：a[u] * a[v] <= k
直接DFS从树根开始搜索，并且使用线段树维护条件2即可。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long LL;

const long long INF = (1LL<<60) - 1;
const int MAX = 1e9;

long long a[200000 + 50], b[200000 + 50];
long long k;
int deep[100000 + 50];
int sum[800000 + 50];
long long ans;
int m, n;

struct line
{
    int to, next;
}edge[200000 + 50];

int pre[100000 + 50];
int tot = 0;
void Add(int u, int v)
{
    edge[tot].to = v;
    edge[tot].next = pre[u];
    pre[u] = tot++;
}

void Build(int left, int right, int root)
{
    if (left == right)
    {
        sum[root] = 0;
        return;
    }
    int mid = (left + right) >> 1;
    Build(left, mid, 2 * root);
    Build(mid + 1, right, 2 * root + 1);
    sum[root] = sum[root * 2] + sum[root * 2 + 1];
}

int Query(int left, int right, int root, int L, int R)
{
    if (L <= left && R >= right)
        return sum[root];
    int mid = (left + right) >> 1;

    if (R <= mid)
        return Query(left, mid, root * 2, L, R);
    else if (L > mid)
        return Query(mid + 1, right, root * 2, L, R);
    else
        return Query(left, mid, root * 2, L, mid) + Query(mid + 1, right, root * 2 + 1, mid + 1, R);
}

void Update(int left, int right, int root, int pos, long long val)
{
    if (left == right)
    {
        sum[root] += val;
        return;
    }
    int mid = (left + right) >> 1;

    if (pos <= mid)
        Update(left, mid, root * 2, pos, val);
    else
        Update(mid + 1, right, root * 2 + 1, pos, val);

    sum[root] = sum[root * 2] + sum[root * 2 + 1];
}

void DFS(int from)
{
    int bound = lower_bound(b + 1, b + m + 1, k / a[from]) - b;
    int pos = lower_bound(b + 1 , b + m + 1, a[from]) - b;
    ans += 1LL*Query(1, m, 1, 1, bound);
    Update(1, m, 1, pos, 1);

    for (int i = pre[from]; i != -1; i = edge[i].next)
        DFS(edge[i].to);
    Update(1, m, 1, pos, -1);
}

int main()
{
    int T;
    scanf("%d", &T);

    while (T--)
    {
        memset(pre, -1, sizeof(pre));
        memset(deep, 0, sizeof(deep));
        memset(sum, 0, sizeof(sum));
        ans = tot = 0;
        scanf("%d %lld", &n, &k);
        for (int i = 1; i <= n; ++i)
        {
            scanf("%lld", &a[i]);
            b[i] = a[i];
        }
        m = n;
        for (int i = 1; i <= n; ++i)
            b[++m] = k/a[i];
        sort(b + 1, b + m + 1);
        m = unique(b + 1, b + m + 1) - (b + 1);
        Build(1, m, 1);
        for (int i = 1; i < n; ++i)
        {
            int u, v;
            scanf("%d %d", &u, &v);
            Add(u, v);
            deep[v]++;
        }
        for (int i = 1; i <= n; ++i)
            if (deep[i] == 0)
                DFS(i);
        printf("%lld\n", ans);
    }
    return 0;
}