黄黑树

from collections import defaultdict
def dfs(root,tmp,node):
    if not node: return 0
    res=0
    if color[root]!=color[node]:
        res+=tmp
    for nxt in dic[node]:
        tmp+=1
        res+=dfs(root,tmp+1,nxt)
        tmp-=1
    return res
if __name__=='__main__':
    n=int(input())
    color=[0]+list(map(int, input().strip().split()))
    fa=[0]+list(map(int, input().strip().split()))
    dic=defaultdict(lambda:[])
    res=[0]*(n+1)
    
    for i in range(1,n):
        dic[fa[i]].append(i+1)
    for i in range(n,0,-1):
        if not dic[i]: continue
        res[i]=int(dfs(i, 0, i)/2)
        
    print(' '.join([str(num) for num in res[1:]]))

#include <iostream>
#include <vector>
#include <map>
#include <queue>
using namespace std;
vector<long long> handle(vector<int>& colors, map<int, vector<int>>& edges) {
    int n = colors.size() - 1;
    vector<long long> ret(n);
    vector<vector<int>> tree;
    queue<int> q;
    q.push(1);
    //构建树并记录每一层的节点
    while(!q.empty()) {
        int sz = q.size();
        vector<int> floor;
        for(int i = 0; i< sz; i++) {
            int node = q.front();
            q.pop();
            floor.push_back(node);
            for(auto j : edges[node]) {
                q.push(j);
            }
        }
        tree.push_back(floor);
    }
    vector<vector<int>> data(n+1, vector<int>(4));
    int m = tree.size();
    // 从最后一层开始遍历，更新每个节点包含的红黑节点数量及其深度之和
    for(int i = m-1; i >= 0; i--) {
        for(auto j : tree[i]) {
            if(edges[j].empty()) {
                if(colors[j]) {
                    data[j] = {0, 0, 1, 0};
                } else {
                    data[j] = {1, 0, 0, 0};
                }
            } else {
                for(auto k : edges[j]){
                    data[j][0] += data[k][0];
                    data[j][1] += data[k][0] + data[k][1];
                    data[j][2] += data[k][2];
                    data[j][3] += data[k][2] + data[k][3];
                }
                if(colors[j]) {
                    data[j][2] += 1;
                    ret[j-1] = data[j][1];
                } else {
                    data[j][0] += 1;
                    ret[j-1] = data[j][3];
                }
            }
        }
    }
    return ret;
}
int main() {
    int n;
    cin >> n;
    vector<int> colors(n+1);
    map<int, vector<int>> edges;
    for(int i = 1; i <= n; i++) {
        cin >> colors[i];
    }
    int node;
    for(int i = 1; i <= n-1; i++) {
        cin >> node;
        edges[node].push_back(i+1);
    }
    for(auto i : handle(colors, edges)) {
        cout << i << " ";
    }
    return 0;
}