题解 | #黄黑树#

典型的递归。根据数据量10^5,如果每个节点都遍历生成结果显然超时，一次遍历实现的话要保证深度相对关系不变。设深度函数f(n,root),节点n对于根节点1的深度为f(n,1)，一个子树包含n,n相对于子树根节点root的深度为f(n,root),root相对于节点1深度为f(root,1),显然有f(n,root)=f(n,1)-f(root,1).即相对深度可由绝对深度之差得到，所以在遍历时只需将根节点的绝对深度减去即可得到相对深度，注意一个同色节点就需减一次，所以在遍历过程中维护四个变量即可：黑点绝对深度和，黑点个数，黄点个数，黄点绝对深度和，同时生成结果。

#include<bits/stdc++.h>
using namespace std;
class Node{
    public:
    int ind;
    int color;
    vector<Node*> sons;
    Node(int i,int c):ind(i),color(c){}
};
class info{
    public:
    int blackSum;
    int blackNum;
    int yellowSum;
    int yellowNum;
    info(int a,int b,int c,int d):blackSum(a),yellowSum(b),blackNum(c),yellowNum(d){}
};

info recur(Node* root,vector<int>& ans,int dep){
    if (root==nullptr) {
       return info(0,0,0,0);
    }
    info cur(0,0,0,0);
    for (auto& son:root->sons){
        info sonInfo=recur(son, ans,  dep+1);
       cur.blackNum+=sonInfo.blackNum;
       cur.blackSum+=sonInfo.blackSum;
        cur.yellowNum+=sonInfo.yellowNum;
        cur.yellowSum+=sonInfo.yellowSum;
    }
   if (root->color==0){
       ans[root->ind]=cur.blackSum-dep*cur.blackNum;
       cur.yellowNum++;
       cur.yellowSum+=dep;
   }else{
       ans[root->ind]=cur.yellowSum-dep*cur.yellowNum;
       cur.blackNum++;
       cur.blackSum+=dep;
   }
    return cur;
}
int main(){
    int n;
    cin>>n;
    vector<Node*> arr(n+1,nullptr);
    int co;
    for (int i=1;i<=n;++i){
        cin>>co;
        arr[i]=new Node(i,co);
    }
    for (int i=1;i<=n-1;++i){
        cin>>co;
        arr[co]->sons.push_back(arr[i+1]);
    }
    vector<int> ans(n+1,0);
    recur(arr[1],ans,0);
    for (int i=1;i<=n;++i) cout<<ans[i]<<" ";
    cout<<endl;
}