题意
给你一棵树,每个节点有一个权值。询问每个节点有多少个点对,满足以该节点为根节点 u u u与它的所有子节点 v v v,并且 ∣ v a l [ u ] − v a l [ v ] ∣ ≤ K \left|val[u] - val[v] \right|\leq K ∣val[u]−val[v]∣≤K
思路
- 比赛的时候想到一个思路:就是每次从叶子节点向根节点返回,每次把经过的点加入集合,并满足根节点的数量。没时间写了,回来实现一下发现超时,就可能树上有一个长链,在底层有很多叶子节点,这样每次相当于便利了一棵树 N 2 N^2 N2的复杂度。
- 再次向学长请教(学长好厉害啊^ _ ^)。
- 按照 d f s dfs dfs顺序维护一个树状数组,在刚进入节点的时候先计算满足根节点的个数 t m p tmp tmp,然后对子节点进行 d f s dfs dfs,在回溯的时候插入当前点再次统计满足根节点的情况 s u m sum sum,那么 s u m − t m p sum - tmp 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;
}