2019 ICPC 南昌 Regional K. Tree（树上启发式合并 + 动态开点线段树）

分析：

对每一个值建立一棵线段树树，保存[l, r]内的深度出现的次数，然后通过启发式合并统计答案即可

代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;//三年竞赛一场空，不开long long见祖宗 
//typedef __int128 lll;
#define print(i) cout << "debug: " << i << endl
#define close() ios::sync_with_stdio(0), cin.tie(0), cout.tie(0)
#define mem(a, b) memset(a, b, sizeof(a))
#define pb(a) push_back(a)
#define x first
#define y second
typedef pair<int, int> par;
const ll mod = 1e9 + 7;
const int maxn = 1e5 + 10;
const int inf = 0x3f3f3f3f;
vector<int> g[maxn];
int val[maxn], siz[maxn], mson[maxn];
int dep[maxn];
int n, k;
ll ans;

void dfs1(int u)
{
    siz[u] = 1;
    for(int v : g[u])
    {
        dep[v] = dep[u] + 1;
        dfs1(v);
        siz[u] += siz[v];
        if(siz[v] > siz[mson[u]]) mson[u] = v;
    }
}
int tot;
int t[maxn];
int lson[maxn * 200], rson[maxn * 200];
int sum[maxn * 200];
void update(int& now, int l, int r, int pos, int val)
{
    if(!now) now = ++tot;
    sum[now] += val;
    if(l == r) return;
    int mid = l + r >> 1;
    if(pos <= mid) update(lson[now], l, mid, pos, val);
    else update(rson[now], mid + 1, r, pos, val);
}

int query(int now, int l, int r, int L, int R)
{
    if(L <= l && r <= R) return sum[now];
    int mid = l + r >> 1;
    ll res = 0;
    if(L <= mid) res += query(lson[now], l, mid, L, R);
    if(R > mid) res += query(rson[now], mid + 1, r, L, R); 
    return res;
}

void update1(int u, int value)
{
    update(t[val[u]], 1, n, dep[u], value);
    for(int v : g[u])
        update1(v, value);
}

void query1(int u, int value, int depth)
{
    int nowval = 2 * value - val[u];
    int maxdep = min(n, 2 * depth + k - dep[u]);
    if(nowval >= 0 && nowval <= n && maxdep >= 1) 
        ans += 2 * query(t[nowval], 1, n, 1, maxdep);
    for(int v : g[u])
        query1(v, value, depth);

}

void dfs2(int u)
{
    for(int v : g[u])
    {
        if(v == mson[u]) continue;
        dfs2(v);
        update1(v, -1);
    }
    if(mson[u]) dfs2(mson[u]);
    for(int v : g[u])
    {
        if(v == mson[u]) continue;
        query1(v, val[u], dep[u]);
        update1(v, 1);  
    }
    update(t[val[u]], 1, n, dep[u], 1);
}

int main()
{
    cin >> n >> k;
    for(int i = 1; i <= n; i++) cin >> val[i];
    for(int i = 2; i <= n; i++)
    {
        int x; cin >> x;
        g[x].pb(i);
    }
    dep[1] = 1;
    dfs1(1);
    dfs2(1);
    cout << ans << endl;
}