https://nanti.jisuanke.com/t/42586

题意：
给定以\(1\)为根的有向树，编号为\(i\)的点具有权值\(v_i\)，问树上存在多少有序对\({x, y}\)，设\(LCA_{x, y} = z\)，使得\(x \neq z\)，\(y \neq z\)，\(x\)和\(y\)的树上距离不超过\(k\)，且\(v_x + v_y = 2 * v_z\)

思路：
树上距离同样可以通过\(LCA\)来求解，那么我们在确定了\(LCA\)之后，找这些点对就比较有想法了。

一次找两个点，同样还是有点困难的，先考虑如何暴力去找。设\(LCA\)为点\(z\)，我们枚举点\(x\)，那么就能算出想要的点\(y\)的深度和权值了。权值是确定的，但是深度是一个区间，考虑一次性对区间进行询问。对于不同的权值，我们每次要询问一个区间，那么就对每个权值建一棵权值线段树，代表某一深度的数量。直接开会太大，采用动态开点的策略。

接下来考虑如何减小时间复杂度。因为\(x\)和\(y\)互相没有祖先关系，如果每次确定点\(x\)所在的链，询问之后再将这条链加入询问集合，就能直接避免出现这种情况，并且再将询问结果\(*2\)就能得到有序对的答案，很像dsu on tree的过程，那么剖完大力合并就行了。

用自己原先的板子很难写这道题，写完顺便更新了板子，学到许多。

#include <bits/stdc++.h>

using namespace std;

#define endl "\n"

typedef long long ll;

const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
const int N = 1e5 + 7;

int n, k, tot, son;
int w[N], dep[N], sz[N], bigSon[N];
int rt[N * 50], ls[N * 80], rs[N * 80], num[N * 80];
ll ans;
vector<int> G[N];

void modify(int &id, int l, int r, int p, int v) {
    if (!id) id = ++tot;
    num[id] += v;
    if (l == r) return;
    int mid = l + r >> 1;
    if (p <= mid) modify(ls[id], l, mid, p, v);
    else modify(rs[id], mid + 1, r, p, v);
}

int query(int id, int l, int r, int ql, int qr) {
    if (l >= ql && r <= qr) return num[id];
    int ans = 0;
    int mid = l + r >> 1;
    if (ql <= mid) ans += query(ls[id], l, mid, ql, qr);
    if (qr > mid) ans += query(rs[id], mid + 1, r, ql, qr);
    return ans;
}

void dfsSz(int u) {
    sz[u] = 1;
    for (auto v : G[u]) {
        dep[v] = dep[u] + 1;
        dfsSz(v);
        sz[u] += sz[v];
        if (sz[v] > sz[ bigSon[u] ]) bigSon[u] = v;
    }
}

void add(int u, int val) {
    modify(rt[ w[u] ], 1, n, dep[u], val);
    for (auto v : G[u]) {
        add(v, val);
    }
}

void calc(int u, int fa) {
    int dv = 2 * dep[fa] + k - dep[u];
    int wv = 2 * w[fa] - w[u];
    dv = min(dv, n);
    if (dv >= 1 && wv >= 0 && wv <= n) ans += 2ll * (ll)query(rt[wv], 1, n, 1, dv);
    for (auto v : G[u]) {
        calc(v, fa);
    }
}

void dfs(int u) {
    for (auto v : G[u]) {
        if (v == bigSon[u]) continue;
        dfs(v);
        add(v, -1);
    }
    if (bigSon[u]) dfs(bigSon[u]);

    for (auto v : G[u]) {
        if (v == bigSon[u]) continue;
        calc(v, u);
        add(v, 1);
    }
    modify(rt[ w[u] ], 1, n, dep[u], 1);
}

void solve() {
    cin >> n >> k;
    for (int i = 1; i <= n; ++i) {
        cin >> w[i];
    }
    for (int i = 2, u; i <= n; ++i) {
        cin >> u;
        G[u].push_back(i);
    }
    dep[1] = 1;
    dfsSz(1);
    dfs(1);
    cout << ans << endl;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);cout.tie(nullptr);
    cout << fixed << setprecision(20);
    int t = 1;
    while (t--) solve();
    return 0;
}