一天,Chika 对大小接近的点对产生了兴趣,她想搞明白这个问题的树上版本,你能帮助她吗?Chika 会给 你一棵有根树,这棵树有 n 个结点,被编号为 1 n,1 号结点是根。每个点有一个权值,i 号结点的权值为 a[i]。如果 u 是 v 的祖先结点,并且 abs(a[u]−a[v]) ≤K,那么 (u,v) 被称作一个“** 大小接近的点对 **”。 对于树上的每个结点 i,你都需要计算以其为根的子树中的“大小接近的点对”的数量。你需要知道:
(1) abs(x) 代表 x 的绝对值。
(2) 每个结点都是其自身的祖先结点.
输入
输入文件的第一行包含两个整数 n (1≤n≤105) 和 k (1≤k≤109),代表树中结点总数, 以及“大小接近的点对”的大小之差的上界。
第二行包含 n 个整数,第 i 个整数是 a[i] (1≤ a[i] ≤109),代表 i 号结点的权值。
第三行包含 n−1 个整数,第 i 个整数是 i+1 号结点的父结点。
输出
输出应该包含n行,每一行包括一个整数。第i行的整数代表以i为根的子树中的“大小接近的点对”的数量。
思路:由于其值比较大,所以我们先对其进行离散化,然后进行DFS序,进入结点时记录一下符合条件的个数,离开时再记录一下,然后更新记录一下。
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int inf = 0x3f3f3f3f;
const int maxn = 1e6 + 10;
int n, k, num;
vector<ll> v[maxn];
ll a[maxn], b[maxn], ls[maxn];
ll c[maxn], Sum[maxn];
ll lowbit(ll i)
{
return i & -i;
}
void update(ll i)
{
while(i <= n)
{
c[i]++;
i += lowbit(i);
}
}
ll sum(ll i)
{
ll res = 0;
while(i > 0)
{
res += c[i];
i -= lowbit(i);
}
return res;
}
void dfs(ll u)
{
ll l = lower_bound(b, b + num, a[u] - k) - b;
ll r = upper_bound(b, b + num, a[u] + k) - b - 1;
ll s1 = sum(r) - sum(l - 1);
update(ls[u]);
ll cost = 0;
for(int i = 0; i < v[u].size(); ++i)
{
ll x = v[u][i];
dfs(x);
cost += Sum[x];
}
ll s2 = sum(r) - sum(l - 1);
Sum[u] = cost + s2 - s1;
}
int main()
{
memset(c, 0, sizeof(c));
memset(Sum, 0, sizeof(Sum));
scanf("%d%d", &n, &k);
for(int i = 1; i <= n; ++i)
{
scanf("%lld", &a[i]);
b[i] = a[i];
}
b[0] = 0; //使下标从1开始
for(int i = 2; i <= n; ++i)
{
ll x;
scanf("%lld", &x);
v[x].push_back(i);
} //下面是离散化
sort(b, b + n + 1);
num = unique(b, b + n + 1) - b;
for(int i = 1; i <= n; ++i)
ls[i] = lower_bound(b, b + num, a[i]) - b;
dfs(1);
for(int i = 1; i <= n; ++i)
printf("%lld\n", Sum[i]);
return 0;
}