Code : Downward Path with Only Last Vertex Black

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

const int white = 0, black = 1;

class Tree {
  public:
    int n;
    vector<vector<int>> adj;
    vector<long long> cnt;
    vector<int> a;

    // cnt[u] is the number of downward paths starting at u when only the
    // last vertex in the path has black color.
    Tree(int n) {
        this->n = n;
        adj.resize(n);
        cnt.resize(n);
        a.resize(n);
    }

    void dfs(int u, int par) {
        cnt[u] = (a[u] == black);
        for (auto child : adj[u]) {
            if (child == par) {
                continue;
            }
            dfs(child, u);

            if (a[u] == white) {
                cnt[u] += cnt[child];
            }
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);
    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }

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

    long long res = 0;
    t.dfs(0, -1);
    for (int i = 0; i < n; i++) {
        res += t.cnt[i];
        cout << t.cnt[i] << " ";
    }
    cout << endl;
    cout << res << endl;
}

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