zzuli 2520: 大小接近的点对

题目链接

题意

给你一棵树，每个节点有一个权值。询问每个节点有多少个点对，满足以该节点为根节点 $u$与它的所有子节点 $v$，并且 $\left|val[u] - val[v] \right|\leq K$

思路

• 比赛的时候想到一个思路：就是每次从叶子节点向根节点返回，每次把经过的点加入集合，并满足根节点的数量。没时间写了，回来实现一下发现超时，就可能树上有一个长链，在底层有很多叶子节点，这样每次相当于便利了一棵树 $N^2$的复杂度。
• 再次向学长请教（学长好厉害啊^ _ ^）。
• 按照 $dfs$顺序维护一个树状数组，在刚进入节点的时候先计算满足根节点的个数 $tmp$，然后对子节点进行 $dfs$，在回溯的时候插入当前点再次统计满足根节点的情况 $sum$，那么 $sum - tmp$就是当前根节点能与子节点形成的点对，同时加上所有子树的和就是当前根节点的答案。
• 为什么在第一次进入节点的时候统计答案呢？
  刚进入节点的时候，树状数组已经插入了一些节点的值，而这些值并不是当前节点的子节点，所以它不会对节点有贡献。
• 权值很大需要离散。

#include <bits/stdc++.h>
#define LL long long
#define P pair<int, int>
#define lowbit(x) (x & -x)
#define mem(a, b) memset(a, b, sizeof(a))
#define mid ((l + r) >> 1)
#define lc rt<<1
#define rc rt<<1|1
using namespace std;
const int maxn = 1e5 + 7;
int tot;
vector<int> g[maxn];
LL a[maxn], b[maxn], c[maxn], ans[maxn], n, k;
void Add(int x) {
	x++;
	while (x <= n) {
		c[x]++;
		x += lowbit(x);
	}
}
LL Sum(int x) {
	x++;
	LL ans = 0;
	while (x) {
		ans += c[x];
		x -= lowbit(x);
	}
	return ans;
}
void dfs(int x) {
	int l, r;
	l = lower_bound(b, b + tot, a[x-1]-k) - b - 1;
	r = upper_bound(b, b + tot, a[x-1]+k) - b - 1;
	ans[x] -= Sum(r) - Sum(l);
	LL sum = 0;
	for (int it : g[x]) {
		dfs(it);
		sum += ans[it];
	}
	int pos = lower_bound(b, b + tot, a[x-1]) - b;
	Add(pos);
	l = lower_bound(b, b + tot, a[x-1]-k) - b - 1;
	r = upper_bound(b, b + tot, a[x-1]+k) - b - 1;
	ans[x] += Sum(r) - Sum(l) + sum;
	return;
}
int main () {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    scanf("%lld %lld", &n, &k);
    for (int i = 0; i < n; ++i) {
    	scanf("%lld", &a[i]);
    	b[i] = a[i];
    }
    sort(b, b + n);
    tot = unique(b, b + n) - b;
    for (int i = 2, x; i <= n; ++i) {
    	scanf("%d", &x);
    	g[x].push_back(i);
    }
    dfs(1);
    for (int i = 1; i <= n; ++i) {
    	printf("%lld\n", ans[i]);
    }
    return 0;
}