题目链接:https://nanti.jisuanke.com/t/42586
题目大意:
思路:我们考虑如果是dis(x, y)==k。直接STL树上启发式合并。map<pair<int, int>, int> map[x]:x节点的子树的深度,和值的节点个数。
如果是<=k。那么查询的深度就是一个范围。我们用轻重链启发式合并。对每个权值建立线段树。下标为深度。维护节点个数和。
#pragma GCC optimize(2)
#include<bits/stdc++.h>
#define LL long long
using namespace std;
inline int read(){
int s=0,w=1;
char ch=getchar();
while(ch<'0'||ch>'9'){if(ch=='-')w=-1;ch=getchar();}
while(ch>='0'&&ch<='9') s=s*10+ch-'0',ch=getchar();
return s*w;
}
struct setTree {
int lc[5000050], rc[5000050], s[5000050], tot=1;
inline void Delete(int n, vector<int> qk){
for(int i=100005; i<=tot; i++){
lc[i]=rc[i]=s[i]=0;
}
//清空用到的根
for(auto x: qk){
lc[x]=rc[x]=s[x]=0;
}
//预留n个节点,作为n个权值的根
tot=n;
}
inline void up_data(int &i, int l, int r, int x, int v){
//cout<<i<<" "<<l<<" "<<r<<" "<<x<<" "<<v<<endl;
if(i==0){
i=++tot;
}
if(l==r){
s[i]+=v;
return ;
}
int mid=(l+r)>>1;
if(x<=mid){
up_data(lc[i], l, mid, x, v);
}
else{
up_data(rc[i], mid+1, r, x, v);
}
s[i]=s[lc[i]]+s[rc[i]];
}
inline int query(int i, int l, int r, int L, int R){
if(i==0||R<L){
return 0;
}
if(l==L&&r==R){
return s[i];
}
int mid=(l+r)>>1;
if(R<=mid) return query(lc[i], l, mid, L, R);
else if(L>mid) return query(rc[i], mid+1, r, L, R);
return query(lc[i], l, mid, L, mid)+query(rc[i], mid+1, r, mid+1, R);
}
}Tree;
int v[100005], w[100005];
int n, k;
LL ans=0;
struct Treedsu{
vector<int> G[100005], qk;
int s[100005], dfn[100005], son[100005], deep[100005], T=0;
inline void dfs(int u, int fa){
s[u]=1; dfn[u]=++T, deep[T]=deep[dfn[fa]]+1; w[T]=v[u];
for(auto x: G[u]){
if(x!=fa){
dfs(x, u); s[u]+=s[x];
son[u]=s[x]>s[son[u]]?x:son[u];
}
}
}
//添加
inline void add(int x){
qk.push_back(w[x]);
Tree.up_data(w[x], 1, 100005, deep[x], 1);
}
//计算贡献
inline LL getans(int u, int x){
LL ans=0;
int s1=2*(w[u]-1)-(w[x]-1)+1;
int r=k+2*deep[u]-deep[x];
r=min(r, 100005);
if(s1>=1&&s1<=100005){
ans+=Tree.query(s1, 1, 100005, 1, r);
}
return ans;
}
inline void Dsu(int u, int fa, int ok){
//把除了重链的所有的子树的答案计算出来
for(auto x: G[u]){
if(x!=fa&&x!=son[u]){
Dsu(x, u, 0);
}
}
//计算u这棵树重链的答案
if(son[u]) Dsu(son[u], u, 1);
for(auto x: G[u]){
if(x!=fa&&x!=son[u]){
//计算轻链的答案
for(int i=0; i<s[x]; i++){
ans+=getans(dfn[u], dfn[x]+i);
}
//合并轻链
for(int i=0; i<s[x]; i++){
add(dfn[x]+i);
}
}
}
add(dfn[u]);//把u自己加进去
//是否保留线段树
if(!ok){
Tree.Delete(100005, qk);
qk.clear();
}
}
}dsu;
int main() {
int x;
scanf("%d%d", &n, &k);
for(int i=1; i<=n; i++){
v[i]=read();
//权值可以为0,线段树节点>0
v[i]++;
}
for(int i=2; i<=n; i++){
x=read();
dsu.G[i].push_back(x);
dsu.G[x].push_back(i);
}
Tree.Delete(100005, dsu.qk);
dsu.dfs(1, 0);
dsu.Dsu(1, 0, 1);//多样例不保留
printf("%lld\n", ans*2);
return 0;
}

京公网安备 11010502036488号