#include "joitour.h"

#include <bits/stdc++.h>

using namespace std;

constexpr int maxn = 2e5 + 5;

vector<int> g[maxn];

int pr[maxn], sz[maxn];

void dfs(int v) {
    sz[v] = 1;
    for (auto u : g[v]) {
        if (u != pr[v]) {
            pr[u] = v;
            dfs(u);
            sz[v] += sz[u];
        }
    }
    for (int i = 0; i < (int)g[v].size(); ++i) {
        if (g[v][i] == pr[v]) {
            swap(g[v][i], g[v].back());
        }
    }
    if (v != 0) {
        g[v].pop_back();
    }
    for (int i = 0; i < (int)g[v].size(); ++i) {
        if (sz[g[v][i]] > sz[g[v][0]]) {
            swap(g[v][i], g[v][0]);
        }
    }
}

int tp[maxn];

int cnt[maxn][3], cnt_heavy[maxn][3];

void calc_cnt(int v) {
    cnt[v][tp[v]]++;
    cnt_heavy[v][tp[v]]++;
    for (auto u : g[v]) {
        calc_cnt(u);
        for (int i = 0; i < 3; ++i) {
            cnt[v][i] += cnt[u][i];
        }
    }
    if (!g[v].empty()) {
        for (int i = 0; i < 3; ++i) {
            cnt_heavy[v][i] += cnt[g[v][0]][i];
        }
    }
}

long long sum[maxn];

long long ans = 0;

void init(int n, vector<int> f, vector<int> u, vector<int> v, int q) {
    for (int i = 0; i < n; ++i) {
        tp[i] = f[i];
    }
    for (int i = 0; i < n - 1; ++i) {
        g[u[i]].push_back(v[i]);
        g[v[i]].push_back(u[i]);
    }
    dfs(0);
    calc_cnt(0);
    ans = cnt[0][0] * 1LL * cnt[0][1] * 1LL * cnt[0][2];
    for (int i = 0; i < n; ++i) {
        if (tp[i] == 1) {
            for (auto j : g[i]) {
                ans -= cnt[j][0] * 1LL * cnt[j][2];
            }
            ans -= (cnt[0][0] - cnt[i][0]) * 1LL * (cnt[0][2] - cnt[i][2]);
        }
    }
}

void change(int X, int Y) {}

long long num_tours() {
    return ans;
}