CF1676G White-Black Balanced Subtrees 题解

题意

给定一棵树，点有黑白之分。

问这棵树有多少个黑白结点数量相同的子树。

思路

令以 $x$ 为根的子树中黑白结点的个数分别为 $s[x][0/1]$。

则 $s[x][0]=\sum s[v][0]+(x\ \text{is black})$，$s[x][1]=\sum s[v][1]+(x\ \text{is white})$（$v$ 是 $x$ 的子结点）。

如果 $s[x][0] = s[x][1]$，那么以 $x$ 为根的子树是平衡的。

代码

#include <cstdio>
#include <cstring>
struct E{int v, t;}e[4050];int T, n, q, c, h[4050], s[4050][2];
void add(int u, int v) {e[++c] = {v, h[u]};h[u] = c;}
int g() {char c = getchar();while(c != 'W' && c != 'B') c = getchar();return c == 'W';}
void dfs(int x)
{
    for(int i = h[x], v;i;i = e[i].t)
        v = e[i].v, dfs(v), s[x][0] += s[v][0], s[x][1] += s[v][1];
    if(s[x][0] == s[x][1]) ++q;
}
int main()
{
    scanf("%d", &T);
    while(T--)
    {
        scanf("%d", &n);c = q = 0;memset(e, 0, sizeof e);
        for(int i = 1;i <= n;++i) h[i] = s[i][0] = s[i][1] = 0;
        for(int i = 2, t;i <= n;++i) scanf("%d", &t), add(t, i);
        for(int i = 1;i <= n;++i) s[i][g()] = 1;
        dfs(1);printf("%d\n", q);
    }
    return 0;
}