CF1676G White-Black Balanced Subtrees

White-Black Balanced Subtrees

题面翻译

给定一棵 $n$ 个结点的树，根结点是 $1$，每个结点是黑色或白色的。

如果一棵树中黑色结点与白色结点数量相同，那么这棵树是“平衡的"。

问这棵 $n$ 个结点的树有多少棵“平衡的”子树。

题目描述

You are given a rooted tree consisting of $n$ vertices numbered from $1$ to $n$ . The root is vertex $1$ . There is also a string $s$ denoting the color of each vertex: if $s_i = \texttt{B}$ , then vertex $i$ is black, and if $s_i = \texttt{W}$ , then vertex $i$ is white.

A subtree of the tree is called balanced if the number of white vertices equals the number of black vertices. Count the number of balanced subtrees.

A tree is a connected undirected graph without cycles. A rooted tree is a tree with a selected vertex, which is called the root. In this problem, all trees have root $1$ .

The tree is specified by an array of parents $a_2, \dots, a_n$ containing $n-1$ numbers: $a_i$ is the parent of the vertex with the number $i$ for all $i = 2, \dots, n$ . The parent of a vertex $u$ is a vertex that is the next vertex on a simple path from $u$ to the root.

The subtree of a vertex $u$ is the set of all vertices that pass through $u$ on a simple path to the root. For example, in the picture below, $7$ is in the subtree of $3$ because the simple path $7 \to 5 \to 3 \to 1$ passes through $3$ . Note that a vertex is included in its subtree, and the subtree of the root is the entire tree.

The picture shows the tree for $n=7$ , $a=[1,1,2,3,3,5]$ , and $s=\texttt{WBBWWBW}$ . The subtree at the vertex $3$ is balanced.

输入格式

The first line of input contains an integer $t$ ( $1 \le t \le 10^4$ ) — the number of test cases.

The first line of each test case contains an integer $n$ ( $2 \le n \le 4000$ ) — the number of vertices in the tree.

The second line of each test case contains $n-1$ integers $a_2, \dots, a_n$ ( $1 \le a_i < i$ ) — the parents of the vertices $2, \dots, n$ .

The third line of each test case contains a string $s$ of length $n$ consisting of the characters $\texttt{B}$ and $\texttt{W}$ — the coloring of the tree.

It is guaranteed that the sum of the values $n$ over all test cases does not exceed $2 \cdot 10^5$ .

输出格式

For each test case, output a single integer — the number of balanced subtrees.

样例 #1

样例输入 #1

3
7
1 1 2 3 3 5
WBBWWBW
2
1
BW
8
1 2 3 4 5 6 7
BWBWBWBW

样例输出 #1

2
1
4

提示

The first test case is pictured in the statement. Only the subtrees at vertices $2$ and $3$ are balanced.

In the second test case, only the subtree at vertex $1$ is balanced.

In the third test case, only the subtrees at vertices $1$ , $3$ , $5$ , and $7$ are balanced.

思路：

树形dp的模板题！我一开始没想好到底要怎么遍历数据，后面仔细看了他们的代码才懂orz。样例的那一组数据输入的是根节点，把i到输入的根节点对应起关系来，建成树。
然后就是在遍历以1为根节点的树时，把每个结点有的白黑结点给求出来，这里在输入WBWB的时候就可以预处理。
f[i][0]表示该结点数下白色结点的个数，f[i][1]表示该结点下黑色结点的个数

代码：

#include<iostream>
#include<cstring>
using namespace std;
const int N=4010;
int h[N],e[N],ne[N],idx;
bool w[N];
int f[N][2];//f[i][0]表示该结点数下白色结点的个数，f[i][1]表示该结点下黑色结点的个数
void add(int a,int b){
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
int res;
void dfs(int u){
    for(int i=h[u];~i;i=ne[i]){
        int j=e[i];
        //遍历它的子节点
        dfs(j);//将它下面的结点统计完成
        f[u][1]+=f[j][1];
        f[u][0]+=f[j][0];
    }
    if(f[u][0]==f[u][1]) {
        //处理完每个结点之后可以计算一下答案
        res++;
    }
}
int main(){
    int t;cin>>t;
    while(t--){
        int n;cin>>n;
        //开始建树
        memset(h,-1,sizeof h);
        memset(e,0,sizeof e);
        memset(ne,0,sizeof ne);
        idx=0;
        memset(w,0,sizeof w);
        memset(f,0,sizeof f);
        for(int i=2;i<=n;i++){
            //因为1是头节点，所以要从2开始建立
            int x;cin>>x;
            add(x,i);//在x和i之间建立一条边，把i插入到x结点下
        }
        for(int i=1;i<=n;i++){
            char c;cin>>c;
            f[i][c=='B']=1;//这个表示如果是黑色结点f[i][1]=1,是白色结点则是f[i][0]=1;
        }
        //开始遍历结点
        res=0;
        dfs(1);
        printf("%d\n",res);
    }
}