树的模板题

#树的模板题
##小早的树
链接：https://ac.nowcoder.com/acm/contest/53417/F
来源：牛客网

小早有一棵树，树上的节点有三种颜色分别为0，1，2，小早希望你统计以某一个节点为根的子树中（不包含这个节点），各个颜色的节点有多少个。
输入描述:
一行两个整数，
1
≤
�
≤
1
0
6
1≤n≤10
6
和
1
≤
�
�
�
�
≤
�
1≤root≤n ，分别表示树的总结点数目和根。
接下来
�
−
1
n−1 行，每行两个正整数
1
≤
�
,
�
≤
�
1≤x,y≤n ，表示节点
�
x 和节点
�
y 之间有连结，保证输入的是一棵树。
接下来一行共
�
n 个正整数，只含有0，1，2，表示第
�
i个节点的颜色。
输出描述:
共计
�
n行，每行三个整数，第
�
i行分别表示以
�
i为根节点的子树（不含
�
i）共有多少个节点颜色为0，1，2
示例1
输入
复制
4 1
1 2
1 3
2 4
0 1 1 1
输出
复制
0 3 0
0 1 0
0 0 0
0 0 0

#include<iostream>
#include<cstring>
#include<algorithm>
using namespace std;
const int N=2000005;
int h[N],e[N],ne[N];
int b[N];
int idx;
bool st[N];
int diyi[N],dier[N],disan[N];
//a所对应的单链表中插入b  a作为根
void add(int a,int b)
{
	e[idx]=b;
	ne[idx]=h[a];
	h[a]=idx++;
}
// dfs 框架
/*
void dfs(int u){
    st[u]=true; // 标记一下，记录为已经被搜索过了，下面进行搜索过程
    for(int i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(!st[j]) {
            dfs(j);
        }
    }
}
*/
void dfs(int u)
{
	st[u]=true;
	for(int i=h[u];i!=-1;i=ne[i]){
		int j=e[i];
		if(!st[j]){
			dfs(j);
			if(b[j]==0){
				diyi[u]+=diyi[j]+1;
				dier[u]+=dier[j];
				disan[u]+=disan[j];
			}
			else if(b[j]==1){
				diyi[u]+=diyi[j];
				dier[u]+=dier[j]+1;
				disan[u]+=disan[j];
			}
			else if(b[j]==2){
				diyi[u]+=diyi[j];
				dier[u]+=dier[j];
				disan[u]+=disan[j]+1;
			}
		}
	}
}
int main()
{
	int n,root;
	cin>>n>>root;
	memset(h,-1,sizeof h);
	for(int i=0;i<n-1;i++)
	{
		int a,b;
		cin>>a>>b;
		add(a,b);
		add(b,a);
	}
	for(int i=1;i<=n;i++)
	{
		cin>>b[i];
	}
	dfs(root);
	for(int i=1;i<=n;i++)
	{
		cout<<diyi[i]<<" "<<dier[i]<<" "<<disan[i]<<endl;
	}
	
	return 0;
}