B. Updating the Tree
time limit per test
1.5 s
memory limit per test
256 MB
input
standard input
output
standard output
A rooted tree is a tree with a countable number of nodes, in which a particular node is distinguished from the others by being the ancestor node of every node of the tree. This node is called the root node.
You are given a rooted tree consisting of n nodes numbered from 1 to n. The root of the tree is node number 1. Each node i has a value viassigned to it.
For each subtree, you must find the minimum number of nodes you must change the value on them to any other value so that the distance between every pair of nodes in that subtree is equal to the absolute difference between the values on them, or say that it is impossible. Can you?
Input
The first line contains an integer T (1 ≤ T ≤ 100) specifying the number of test cases.
The first line of each test case contains an integer n (1 ≤ n ≤ 105), in which n is the number of nodes in the tree. Then a line follow containing n integers v1, ..., vn (1 ≤ vi ≤ 105), in which vi is the value of the ith node.
Then n - 1 lines follow, each line contains two integers ai and bi (1 ≤ ai, bi ≤ n), giving an edge between nodes ai and bi.
Output
For each test case, print a single line containing n space-separated integers x1, ..., xn, in which xi is the answer of the subtree of node i. If an answer does not exist for a subtree, print - 1.
Example
input
Copy
1 2 1 3 1 2
output
Copy
1 0
解题思路：
如果一颗子树是三叉树，那么当前点肯定没答案，然后他的所有父亲都肯定没答案。
如果一颗子树是二叉树，那么当前点有答案，但是他的所有父亲都没有答案。
所以我们要把所有的链都取出来，然后如果是二叉树，就要把他两个孩子的两条链合并。
先考虑一条链的情况。假设i的深度大于j，那么题目就是要使得
| ai - aj | = depi - depj
假设 ai >= aj
那么就是要求
ai - depi = aj - depj
假设 ai < aj
那么就是
ai + depi = aj + depj
所以我们对于一条链而言，维护每个点的差和和的数量，用子树大小减去最大的数量就是答案。
如果是二叉树，就要合并孩子的那两条链。
合并的时候用 左孩子的和跟右孩子的差合并，差跟和合并，同时要加上或减去偏移量
| ai - aj | = depi + depj - 2*dep[LCA(i,j)]
所以我们用一遍深搜，先求出肯能有答案的节点。
然后第二遍深搜，深搜到可能有答案的节点，然后从那个节点开始深搜，然后从最底下回溯，如果是二叉树，就处理下即可。
这里要注意根节点的特殊判断。
详见代码
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN = 200005;
int a[MAXN];
vector<int> G[MAXN];
bool vis[MAXN];//记录是否可能有答案呢
bool vis2[MAXN];//记录是否已经计算出答案
int dep[MAXN];//深度
int sz[MAXN];//子树大小
//判断是否能有答案
void dfs1(int u,int fa,int d){
dep[u]=d;
sz[u]=1;
if(G[u].size()>3)
vis[u]=1;
for(auto &v:G[u]){
if(v==fa)
continue;
dfs1(v,u,d+1);
sz[u]+=sz[v];
if(G[v].size()>=3||vis[v])//孩子没答案，父亲也肯定没答案
vis[u]=1;
}
}
map<int,int> add[2];//合并用
map<int,int> sub[2];
int maxx=0;
int ans[MAXN];
int combine(int u){
for(auto &t:sub[0]){
add[1][t.first+2*dep[u]]+=t.second;//合并记得加上偏移量
maxx=max(maxx,add[1][t.first+2*dep[u]]);
}
for(auto &t:add[0]){
sub[1][t.first-2*dep[u]]+=t.second;
maxx=max(maxx,sub[1][t.first-2*dep[u]]);
}
add[1][a[u]+dep[u]]++;
sub[1][a[u]-dep[u]]++;
maxx=max(maxx,add[1][a[u]+dep[u]]);
maxx=max(maxx,sub[1][a[u]-dep[u]]);
return sz[u]-maxx;
}
//处理一条链的情况
void dfs3(int u,int fa,int c){
vis2[u]=1;
if(u!=1)
{
if(G[u].size()==1)//到了根节点，开始回溯
{
add[c][a[u]+dep[u]]++;
sub[c][a[u]-dep[u]]++;
maxx=max(maxx,add[c][a[u]+dep[u]]);
maxx=max(maxx,sub[c][a[u]-dep[u]]);
ans[u]=sz[u]-maxx;
}
}
for(auto &v:G[u]){
if(v==fa)
continue;
dfs3(v,u,c);
add[c][a[u]+dep[u]]++;
sub[c][a[u]-dep[u]]++;
maxx=max(maxx,add[c][a[u]+dep[u]]);//回溯统计答案
maxx=max(maxx,sub[c][a[u]-dep[u]]);
ans[u]=sz[u]-maxx;
}
}
void dfs2(int u,int fa){
if(vis[u]==0){
if(G[u].size()<=2&&u!=1)//一条链的情况
{
maxx=0;
dfs3(u,fa,0);
add[0].clear();
sub[0].clear();
}
else{
//二叉树的情况
int cnt=0;
for(auto &v:G[u]){
if(v==fa)
continue;
maxx=0;
dfs3(v,u,cnt++);//先深搜两个孩子（肯定都是链）
}
ans[u]=combine(u);//然后合并
add[0].clear();
add[1].clear();
sub[0].clear();
sub[1].clear();
}
}
for(auto &v:G[u]){
if(v==fa)
continue;
if(vis2[v])
continue;
dfs2(v,u);
}
}
int main()
{
int T;
scanf("%d",&T);
while(T--){
memset(ans,-1,sizeof(ans));
memset(vis,0,sizeof(vis));
memset(vis2,0,sizeof(vis2));
add[0].clear();
add[1].clear();
sub[0].clear();
sub[1].clear();
int N;
scanf("%d",&N);
for(int i=1;i<=N;i++){
scanf("%d",&a[i]);
G[i].clear();
}
for(int i=1,u,v;i<N;i++){
scanf("%d%d",&u,&v);
G[u].push_back(v);
G[v].push_back(u);
}
if(N==1)//特殊处理
{
printf("%d\n",0);
}
else{
if(G[1].size()>2)//特殊处理
vis[1]=1;
dfs1(1,0,0);
dfs2(1,0);
for(int i=1;i<=N;i++){
printf("%d%c",ans[i],i==N?'\n':' ');
}
}
}
return 0;
}