Weak Pair （HDU - 5877 ，DFS + 离散化 + 权值线段树）

一.题目链接：

HDU-5877

二.题目大意：

一棵树，每个节点有一个权值 a[i]。

一个 weak对(u, v) 需满足：

① u 是 v 的祖先

② 

问有多少个 weak 对.

三.分析：

dfs 一棵树，对于当前节点 v 来说，只需找权值 ≤ k / a[v] 的祖先有多少个.

权值线段树查询，每经过一个点便将 a[u] 的次数 + 1，结束时将 a[u] 的次数 - 1.

由于 k 过大，这里还需要离散化.

详见代码.

四.代码实现：

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int M = (int)1e5;
const ll inf = 0x3f3f3f3f3f3f3f3f;

int n, m;
int a[M + 5];
bool in[M + 5];
ll ans, k, b[M * 2 + 5];

int cnt;
int head[M + 5];
struct enode
{
    int v, nx;
}Edge[M + 5];

int tree[M * 8 + 5];

void init(int n)
{
    m = ans = cnt = 0;
    for(int i = 1; i <= n; ++i)
    {
        in[i] = 0;
        head[i] = -1;
    }
}

void add(int u, int v)
{
    Edge[cnt].v = v;
    Edge[cnt].nx = head[u];
    head[u] = cnt++;
}

inline int lc(int k)
{
    return k<<1;
}

inline int rc(int k)
{
    return k<<1|1;
}

void build(int k, int l, int r)
{
    tree[k] = 0;
    if(l == r)
        return;
    int mid = (l + r) >> 1;
    build(lc(k), l, mid);
    build(rc(k), mid + 1, r);
}

int query(int k, int l, int r, int a, int b)
{
    if(l >= a && r <= b)
        return tree[k];
    int mid = (l + r) >> 1;
    int sum = 0;
    if(a <= mid)
        sum += query(lc(k), l, mid, a, b);
    if(mid < b)
        sum += query(rc(k), mid + 1, r, a, b);
    return sum;
}

void update(int k, int l, int r, int pos, int v)
{
    if(l == r)
    {
        tree[k] += v;
        return;
    }
    int mid = (l + r) >> 1;
    if(pos <= mid)
        update(lc(k), l, mid, pos, v);
    else
        update(rc(k), mid + 1, r, pos, v);
    tree[k] = tree[lc(k)] + tree[rc(k)];
}

void dfs(int u)
{
    int r = lower_bound(b + 1, b + m + 1, a[u] ? k / a[u] : inf) - b;
    ans += query(1, 1, m, 1, r);
    int pos = lower_bound(b + 1, b + m + 1, a[u]) - b;
    update(1, 1, m, pos, 1);
    for(int i = head[u]; ~i; i = Edge[i].nx)
        dfs(Edge[i].v);
    update(1, 1, m, pos, -1);
}

int main()
{
    int T;
    scanf("%d", &T);
    while(T--)
    {
        scanf("%d %lld", &n, &k);
        init(n);
        for(int i = 1; i <= n; ++i)
        {
            scanf("%d", &a[i]);
            b[++m] = a[i], b[++m] = (a[i] ? k / a[i] : inf);
        }
        sort(b + 1, b + m + 1), m = unique(b + 1, b + m + 1) - (b + 1);
        build(1, 1, m);
        for(int i = 1, u, v; i < n; ++i)
        {
            scanf("%d %d", &u, &v);
            in[v] = 1;
            add(u, v);
        }
        for(int i = 1; i <= n; ++i)
        {
            if(!in[i])
            {
                dfs(i);
                break;
            }
        }
        printf("%lld\n", ans);
    }
    return 0;
}