2019 ICPC 南昌 Regional K. Tree（樹上啓發式合併 + 動態開點線段樹）

題意

求以$1$爲根的$n$個點的有根樹上，滿足下面條件的有序點對$(x,y)$數目（點$i$的權值爲$v_i$）:

1. 一個點不爲另一個點的祖先
2. $v_x+v_y=v_{lca(x,y)}$
3. $x$到$y$的路徑長度小於等於給定的值 $k$

$n,k\le 1e5, 0\le v_i\le n$

解題思路：

樹上啓發式合併，先統計輕兒子子樹中的答案，對輕兒子處理完之後清除輕兒子的痕跡，然後統計重兒子子樹中的答案。統計完之後保留重兒子的痕跡，依次遍歷輕兒子的結點去詢問答案，處理完一顆輕子樹之後把這顆子樹合並進來。
具體的，對每一個權值$x$開一顆線段樹，存放權值爲$x$，深度在$[l,r]$範圍的點有多少個，並以此爲根據來詢問。

自己踩的坑：
be up to k理解爲必須爲k，人傻了
算出來的深度範圍如果上限超過n，應該以n爲上限，而不是直接不詢問（上一個坑的後遺症）
所以英語很重要

#include<bits/stdc++.h>
#define ll long long
#define mid ((l+r)>>1)
using namespace std;
const int maxn = 1e5 + 50;
int n, k;
vector<int> g[maxn];
int val[maxn], son[maxn], dep[maxn], sz[maxn];
int T[maxn], lc[maxn*200], rc[maxn*200], tot = 0, sum[maxn*200];
void update(int &rt, int l, int r, int pos, int x){
    if(!rt) rt = ++tot;
    sum[rt] += x;
    if(l == r) return;
    if(pos <= mid) update(lc[rt], l, mid, pos, x);
    else update(rc[rt], mid+1, r, pos, x);
}
int qry(int rt, int l, int r, int L, int R){
    if(!rt) return 0;
    if(L <= l && r <= R) return sum[rt];
    int res = 0;
    if(L <= mid) res += qry(lc[rt], l, mid, L, R);
    if(R > mid) res += qry(rc[rt], mid+1, r, L, R);
    return res;
}
ll ans = 0;
void init(){
    scanf("%d%d", &n, &k);
    for(int i = 1; i <= n; ++i) scanf("%d", &val[i]);
    for(int i = 2; i <= n; ++i){int u; scanf("%d", &u); g[u].push_back(i);}
}
void dfs1(int u){
    sz[u] = 1;
    for(int i = 0; i < g[u].size(); ++i){
        int v = g[u][i];
        dep[v] = dep[u]+1;
        dfs1(v);  sz[u] += sz[v];
        if(sz[v] > sz[son[u]]) son[u] = v;
    }return;
}
void del(int u){
    update(T[val[u]], 1, n, dep[u], -1);
    for(int i = 0; i < g[u].size(); ++i) del(g[u][i]);
}
void qry(int u, int td, int tv){
    int d = k+2*td-dep[u];
    d = min(d, n);//Attention!!
    int t = 2*tv-val[u];
    if(d >= 1 && t >= 0 && t <= n)  ans = ans + 2LL*qry(T[t], 1, n, 1, d);
    for(int i = 0; i < g[u].size(); ++i) qry(g[u][i], td, tv);
}
void add(int u){
    update(T[val[u]], 1, n, dep[u], 1);
    for(int i = 0; i < g[u].size(); ++i) add(g[u][i]);
}
void dfs2(int u){

    for(int i = 0; i < g[u].size(); ++i){
        int v = g[u][i];
        if(v == son[u]) continue;
        dfs2(v);del(v);
    }
    if(son[u]) dfs2(son[u]);
    for(int i = 0; i < g[u].size(); ++i){
        int v = g[u][i];  if(v == son[u]) continue;
        qry(v, dep[u], val[u]);  add(v);
    }
    update(T[val[u]], 1, n, dep[u], 1);
}
int main()
{
    init();
    dep[1] = 1;
    dfs1(1);dfs2(1);
    cout<<ans<<endl;
}