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

传送门

题意

分析

首先，要求一个点不为另一个点的祖先，这个比较好处理，先确定一个根结点$u$，在他的两颗子树里面找点就可以了
两个点和权值和为定值，可以用离散化+树上启发式合并进行处理
剩下的就是路径长度了，这个不太好处理，因为是求一个范围内的合法点的数量，怎么办呢，我们可以用权值线段树来处理，每一个点权建一颗线段树，表示每一个深度的点出现了多少次，因为范围比较大，所以需要动态开点

代码

#pragma GCC optimize(3)
#include <bits/stdc++.h>
#define debug(x) cout<<#x<<":"<<x<<endl;
#define dl(x) printf("%lld\n",x);
#define di(x) printf("%d\n",x);
#define _CRT_SECURE_NO_WARNINGS
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(),(x).end()
#define fi first
#define se second
#define SZ(x) ((int)(x).size())
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> PII;
typedef vector<int> VI;
const int INF = 0x3f3f3f3f;
const int N = 1e5 + 50;
const ll mod = 1000000007;
const double eps = 1e-9;
const double PI = acos(-1);
template<typename T>inline void read(T &a) {
	char c = getchar(); T x = 0, f = 1; while (!isdigit(c)) {if (c == '-')f = -1; c = getchar();}
	while (isdigit(c)) {x = (x << 1) + (x << 3) + c - '0'; c = getchar();} a = f * x;
}
int gcd(int a, int b) {return (b > 0) ? gcd(b, a % b) : a;}
VI G[N];
int a[N],sz[N],son[N],dep[N];
int Son,idx;
ll ans;
int root[N];
int n,k;
struct Node{
	int l,r;
	int x;	 
}tr[N * 200];

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

int qeruy(int u,int l,int r,int L,int R){
	if(!u) return 0;
	if(l >= L && r <= R) return tr[u].x;
	int mid = (l + r) >> 1;
	int x = 0;
	if(L <= mid) x += qeruy(tr[u].l,l,mid,L,R);
	if(R > mid) x += qeruy(tr[u].r,mid + 1,r,L,R);
	return x; 
}

void modify(int &u,int l,int r,int p,int x){
	if(!u) u = ++idx;
	tr[u].x += x;
	if(l == r) return;
	int mid = (l + r) >> 1;
	if(p <= mid) modify(tr[u].l,l,mid,p,x);
	else modify(tr[u].r,mid + 1,r,p,x);
}

void run(int u,int td,int tv){
	int d = k + 2 * td - dep[u];
	d = min(d,n);
	int t = tv * 2 - a[u];
	if(d >= 1 && t >= 0 && t <= n) ans = (ans + 2ll * qeruy(root[t],1,n,1,d));
	for(auto v:G[u]) run(v,td,tv);
}

void add(int u,int x){
	modify(root[a[u]],1,n,dep[u],x);
	for(auto v:G[u]) add(v,x);
}

void dsu(int u){
	for(auto v:G[u]){
		if(v == son[u]) continue;
		dsu(v);
		add(v,-1);
	}
	if(son[u]) dsu(son[u]);
	for(auto v:G[u]){
		if(v == son[u]) continue;
		run(v,dep[u],a[u]);
		add(v,1);
	}	
	modify(root[a[u]],1,n,dep[u],1);

}	

int main() {
	read(n),read(k);
	dep[1] = 1;
	for(int i = 1;i <= n;i++) read(a[i]);
	for(int i = 2;i <= n;i++){
		int x;
		read(x);
		G[x].pb(i);
	}
	dfs(1);
	dsu(1);
	dl(ans);
	return 0;
}