Codeforces Round #756 (Div. 3) D - Weights Assignment For Tree Edges

D - Weights Assignment For Tree Edges

题意
给定b[]，含义是i的父节点是b[i];给定p[]，要求最后构建出来的树，节点到root的距离从小到大，排列顺序与p[]的顺序相同。
思路：
定义两个数组，

1. dp[maxn]：i到root的距离为dp[i] ,
2. ans[maxn]：i到其父节点的距离为ans[i]，

最后要输出的答案就是ans数组。从1到n遍历p[]，维护更新dp[]和ans[]即可。

1. ans[p[i]]=(dp[p[i-1]]+x)-dp[b[p[i]]];
   p[i]到其父节点的距离 = 上一个节点到root距离+x - p[i]的父节点到root的距离
   上一个点是指在数组p中i的上一个点即p[i-1]
   此处x>=1即可，只要保证距离是按p[]递增即可

2. dp[p[i]]=dp[b[p[i]]]+ans[p[i]];
   p[i]到root距离 = p[i]的父节点到root距离 + p[i]到其父节点的距离

#include<bits/stdc++.h>
using namespace std;
const int maxn=2e5+5;

int b[maxn],p[maxn];
int dp[maxn];//i到root的距离 
int ans[maxn];//i到其父节点的距离 
int n,root;

int main(){
	ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
	
	int t;
	cin>>t;
	while(t--){
		
		cin>>n;
		for(int i=1;i<=n;i++){
			cin>>b[i];
			if(b[i]==i)	root=i;
		}
		for(int i=1;i<=n;i++)	cin>>p[i];
		
		bool flag=true;
		if(p[1]!=root)	flag=false;
		for(int i=2;i<=n;i++){//此循环可以在cin>>p[i]时进行，但是最好是分开处理，血的教训 
			if(b[p[i]]!=root&&dp[b[p[i]]]==0){//如果p[i]节点的父节点还未处理，那么一定-1 
				flag=false;
				break;
			}
			ans[p[i]]=(dp[p[i-1]]+1)-dp[b[p[i]]];//p[i]到其父节点的距离 = 上一个节点到root距离+1 - p[i]的父节点到root的距离 
			dp[p[i]]=dp[b[p[i]]]+ans[p[i]];//p[i]到root距离 = p[i]的父节点到root距离 + p[i]到其父节点的距离
		}
		
		if(flag){
			for(int i=1;i<=n;i++)	cout<<ans[i]<<" ";
		}else	cout<<"-1";
		cout<<endl;
		
		for(int i=1;i<=n;i++)	ans[i]=dp[i]=0;
	}
	return 0;
}
/*
4
5
3 1 3 3 1
3 1 2 5 4
3
1 1 2
3 1 2
7
1 1 2 3 4 5 6
1 2 3 4 5 6 7
6
4 4 4 4 1 1
4 2 1 5 6 3
*/