小团作为一名美团骑手,最喜欢的颜色就是黄和黑,因此对包含这两种颜色的事物都格外留意。
这天他发现有一棵树,树上的每个节点都是黄的或者黑的。现在小团想知道对于这棵树中的每个节点,在以其为根的子树中,所有与其颜色不同的节点的深度之和是多少。子树中节点的深度被定义为该节点与该子树根节点之间的最短路径的边数。
小团作为一名美团骑手,最喜欢的颜色就是黄和黑,因此对包含这两种颜色的事物都格外留意。
这天他发现有一棵树,树上的每个节点都是黄的或者黑的。现在小团想知道对于这棵树中的每个节点,在以其为根的子树中,所有与其颜色不同的节点的深度之和是多少。子树中节点的深度被定义为该节点与该子树根节点之间的最短路径的边数。
第一行有一个数n,代表这棵树上一共有n个节点,编号为1~n。(1<n<=100000)
第二行有n个数,第i个数表示编号为i的节点的颜色,为0表示黄色,为1表示黑色。
第三行有n-1个数,第i个数表示编号为i+1的节点的父节点编号。
在一行中输出n个数,第i个数代表第i个节点的答案。
10 0 0 1 0 0 1 1 1 0 0 1 2 3 4 4 5 7 6 9
17 13 10 6 3 3 0 0 0 0
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; }