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

题目链接

题意

给定一棵树，点上有权值。问多少对点 $(u,v)$ 满足 $u$ 是 $v$ 的祖先 且 $val[u]*val[v]\leq k$.

思路

类似dfs序 题目小集-hdu 3887

注意点

1. 因为 $k\leq 1e18, val\leq 1e9$，所以需离散化，离散化的时候可以将 $val[i]$ 及 $k/val[i]$ 均放入新数组。
   （一开始做的时候每次用到 $k/val[i]$ 都进去 $lower\_bound$ 一下我也是心大…。）
2. 注意到 $val[i]$ 可能为 $0$，此时赋 $inf$.
3. 树状数组的上限不是 $n$ 而是 $q$！！！(q = unique(a, a+q) - a;)
   （因为这个 $WA$ 了很久…一开始 $query$ 还忘了 return ret;…）

Code

#include <bits/stdc++.h>
#define maxn 200010
#define inf 0x3f3f3f3f3f3f3f3f
using namespace std;
typedef long long LL;
LL c[maxn], val[maxn], k, a[maxn], ans;
int ne[maxn], tot, cnt, n, le[maxn], ri[maxn], rev[maxn], inv[maxn], q, ind[maxn];
map<int, set<int> > m;
map<LL, int> nval;
struct Edge {
    int to, ne;
    Edge(int a = 0, int b = 0) : to(a), ne(b) {}
}edge[maxn];
struct Ans { int pre, suc; }rec[maxn];
void addEdge(int u, int v) {
    edge[tot] = Edge(v, ne[u]);
    ne[u] = tot++;
}
int lowbit(int x) { return x & -x; }
void add(int x, LL d) { while (x <= q) c[x] += d, x += lowbit(x); }
LL query(int x) { LL ret = 0; while (x) ret += c[x], x -= lowbit(x); return ret; }
void dfs(int u, int fa) {
    add(lower_bound(a, a+q, val[u]) - a + 1, 1);
    LL temp = query(lower_bound(a, a+q, val[u] ? k / val[u] : inf) - a + 1);
    for (int i = ne[u]; ~i; i = edge[i].ne) {
        int v = edge[i].to;
        if (v == fa) continue;
        dfs(v, u);
    }
    ans += query(lower_bound(a, a+q, val[u] ? k/val[u] : inf) - a + 1) - temp;
}
void work() {
    tot = ans = 0;
    memset(ne, -1, sizeof(ne));
    memset(ind, 0, sizeof(ind));
    memset(c, 0, sizeof(c));

    scanf("%d%lld", &n, &k);
    q = 0;
    for (int i = 1; i <= n; ++i) {
        scanf("%lld", &val[i]);
        a[q++] = val[i], a[q++] = val[i] ? k / val[i] : inf;
    }
    sort(a, a+q);
    q = unique(a, a + q) - a;

    for (int i = 1; i < n; ++i) {
        int u, v;
        scanf("%d%d", &u, &v);
        addEdge(u, v);
        ++ind[v];
    }
    int root;
    for (int i = 1; i <= n; ++i) if (!ind[i]) { root = i; break; }
    dfs(root, -1);
    printf("%lld\n", ans);
}
int main() {
    int T;
    scanf("%d", &T);
    while (T--) work();
    return 0;
}