Editorialist: Tommoa

Congrats to 1021839 (Albert), qazokm (Nick O) and Missile (Alex) for being the only people that managed to get the full marks for this problem in the contest!

The concept for this problem is relatively simple, you need to find the number of ways you can have a colourful tree. Something that would help is noting that the maximum number of unique ways we can colour a tree with \(N\) nodes is (given the tree is a line and all the nodes are '?'s) \(3 \times 2^N\). This will be helpful for complexity analysis for each of the sub-problems.

We had two sub-problems that were increasingly difficult, so lets cover them in order.

20 Points

For 20 points, you had up to \(6\) nodes. Our sample space is small enough that you could brute force the solution relatively easily as there are at most \(3 \times 2^5 = 96\) unique ways to colour the tree in the worst case brute force of a line of '?'s.

80 Points

For the 80 point sub-problem, and eternal bragging rights, we had up to \(N \leq 100\). Using our formula from before, the maximum number of unique ways of colouring the tree is roughly \(2000000000000000000000000000000\) ways (\(3 \times 2^{99} \approx 1.9 \times 10^{30}\)). Clearly a brute force will not work here.

Lets think about this problem a little more. For each node, we select a colour for its children and then calculate the number of ways we can colour the subtree starting at that child. Do we see a reoccurring pattern here? (Hint: the answer is yes).

For each subtree beginning at a node of a colour, we only need to calculate its answer once and then we know what the answer will be for that subtree when the starting node of the same colour is asked for again. This is a concept known as memoization, something relatively common in dynamic programming.

Now instead of having to search for each unique tree, we only have to calculate 3 times for each subtree (of which there are \(N\)), reducing the number of times we need to calculate each node!

An example solution of this is as follows:

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1e9+7;
int T, N;
int memo[300], cols[100]; // 3*100 is the max number of results we'll need
vector<vector<int>> children; // What is what's child?

int rec(int node, int colour) { // colour is the colour of this node
    if (cols[node] != -1 && cols[node] != colour)
        return 0;
    if (memo[colour + (node*3)] != -1) // Have we already calculated this?
        return memo[colour + node*3];
    long long ways = 1;
    for (auto &child: children[node]) {
        long long child_ways = 0;
        for (int col = 0; col < 3; ++col) {
            if (col != colour) { // Make sure we don't have a child with the same colour as the parent
                child_ways += rec(child, col);
                child_ways %= MOD;
            }
        }
        ways *= child_ways;
        ways %= MOD;
    }
    memo[colour + (node*3)] = ways; // Store our result
    return ways;
}

int main() {
    cin >> T;
    while (T--) {
        cin >> N;
        // Reset our arrays
        fill_n(memo, 300, -1);
        fill_n(cols, 100, -1);
        children.assign(N, vector<int>());
        for (int i = 0; i < N; ++i) {
            char colour;
            cin >> colour;
            if (colour == 'C') {
                cols[i] = 0;
            } else if (colour == 'R') {
                cols[i] = 1;
            } else if (colour == 'Y') {
                cols[i] = 2;
            }
        }
        for (int i = 1; i < N; ++i) {
            int parent;
            cin >> parent;
            children[parent].push_back(i);
        }
        long long ways = 0;
        for (int col = 0; col < 3; ++col) {
            ways += rec(0, col);
            ways %= MOD;
        }
        cout << ways << endl;
    }
}