Code : Iris and Game on the Tree

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

const int blue = 0, green = 1, red = 2;
const int middle = 3;

void solve() {
    int n;
    cin >> n;

    vector<vector<int>> adj(n);
    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--, y--;
        adj[x].push_back(y);
        adj[y].push_back(x);
    }

    string str;
    cin >> str;

    auto color = [](char ch) {
        if (ch == '0') {
            return blue;
        };
        if (ch == '1') {
            return green;
        }
        return red;
    };

    map<int, int> count;
    for (int i = 1; i < n; i++) {
        // Leaf.
        if (adj[i].size() == 1) {
            count[color(str[i])]++;
        } else {
            if (color(str[i]) == red) {
                count[middle]++;
            }
        }
    }

    char root = str[0];

    long long ans = 0;
    if (color(root) == green) {
        ans = count[blue];
        ans += (count[red] + 1) / 2;
    }
    if (color(root) == blue) {
        ans = count[green];
        ans += (count[red] + 1) / 2;
    }
    if (color(root) == red) {
        ans = max(count[green], count[blue]);
        // If count of non terminal is even, opponent will just mirror your
        // move.
        if (count[green] != count[blue] || (count[middle] % 2 == 0)) {
            ans += count[red] / 2;
        } else {
            ans += (count[red] + 1) / 2;
        }
    }
    cout << ans << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
}