White-Black Balanced Subtrees (Codeforces 1676G)
Tags: DP, trees

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. 

Input: 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$$$.

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

Note: 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.

Sample Inputs:
3
7
1 1 2 3 3 5
WBBWWBW
2
1
BW
8
1 2 3 4 5 6 7
BWBWBWBW

Sample Outputs:
2
1
4

Solution (Python):
import math
import itertools
import heapq
from math import floor, ceil
from collections import deque
import sys
from collections import Counter
I = lambda : list(map(int, input().split()))
M = lambda : map(int, input().split())
input = sys.stdin.readline
sys.setrecursionlimit(10000)
MOD = 10 ** 9 + 7

def lcm(a, b):
    return abs(a * b) // math.gcd(a, b)

def gcd(a, b):
    if b == 0:
        return a
    return gcd(b, a % b)

def solve(total):
    ans = []
    (m4, m6) = (total // 4, total // 6)
    for i in range(0, m4 + 1):
        for j in range(0, m6 + 1):
            if 4 * i + 6 * j == total:
                ans.append(i + j)
    return ans

def isSubsequence(s, t):
    for i in range(0, len(s)):
        try:
            index = t.index(s[i])
        except ValueError:
            return False
        t = t[index + 1:]
    return True

def bisect_right(a, x):
    (lo, hi) = (0, len(a))
    while lo < hi:
        mid = lo + (hi - lo) // 2
        if a[mid] > x:
            hi = mid
        else:
            lo = mid + 1
    return lo

def bisect_left(a, x):
    (lo, hi) = (0, len(a))
    while lo < hi:
        mid = lo + (hi - lo) // 2
        if a[mid] < x:
            lo = mid + 1
        else:
            hi = mid
    return lo

def binary_search(nums, target):
    (left, right) = (0, len(nums) - 1)
    while left <= right:
        pivot = left + (right - left) // 2
        if nums[pivot] == target:
            return pivot
        if target < nums[pivot]:
            right = pivot - 1
        else:
            left = pivot + 1
    return -1

def fun(l1, l2):
    return l1

def solve():
    N = int(input())
    arr = I()
    col = input()
    DP = [0] * N
    for i in range(N - 1, -1, -1):
        if col[i] == 'W':
            DP[i] += 1
        else:
            DP[i] -= 1
        if i > 0:
            DP[arr[i - 1] - 1] += DP[i]
    print(DP.count(0))
TC = 1
TC = int(input())
for testcases in range(TC):
    solve()

Solution (Python):
for f in range(int(input())):
    n = int(input())
    p = [*map(int, input().split())]
    s = input()
    d = [0] * n
    for i in range(n - 1, -1, -1):
        d[i] += 1 if s[i] == 'W' else -1
        if i > 0:
            d[p[i - 1] - 1] += d[i]
    print(d.count(0))

Solution (Python):
n = int(input())
for i in range(n):
    l = int(input())
    a = list(map(int, input().split()))
    s = input()
    b = [0] * len(s)
    for j in range(len(s) - 1, 0, -1):
        if s[j] == 'W':
            b[j] += 1
            b[a[j - 1] - 1] += b[j]
        else:
            b[j] -= 1
            b[a[j - 1] - 1] += b[j]
    if s[0] == 'W':
        b[0] += 1
    else:
        b[0] -= 1
    print(b.count(0))