#include "joitour.h"

#include <bits/stdc++.h>
#define int long long

using namespace std;

constexpr int R = 1 << 18;

struct Node {
    array<int, 3> cnt;
    array<long long, 3> sum_cnt;
    array<long long, 3> sum_heavy;

    Node() {
        cnt.fill(0);
        sum_cnt.fill(0);
        sum_heavy.fill(0);
    }

    Node operator+(Node b) {
        Node ans;
        for (int i = 0; i < 3; ++i) {
            ans.cnt[i] = cnt[i] + b.cnt[i];
            ans.sum_cnt[i] = sum_cnt[i] + b.sum_cnt[i];
            ans.sum_heavy[i] = sum_heavy[i] + b.sum_heavy[i];
            // if (cnt[1]) {
            //     ans.sum_cnt[i] += sum_cnt[i];
            //     ans.sum_heavy[i] += sum_heavy[i];
            // }
            // if (b.cnt[1]) {
            //     ans.sum_cnt[i] += b.sum_cnt[i];
            //     ans.sum_heavy[i] += b.sum_heavy[i];
            // }
        }
        return ans;
    }

    // void add(Node b) {
    //     for (int i = 0; i < 3; ++i) {
    //         cnt[i] += b.cnt[i];
    //         sum_cnt[i] += b.sum_cnt[i];
    //         sum_heavy[i] += b.sum_heavy[i];
    //     }
    // }
};

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_all[3];

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

int t = 0;

int tin[maxn], tout[maxn];

int ord[maxn];

void dfs2(int v) {
    tin[v] = t++;
    ord[tin[v]] = v;
    for (auto u : g[v]) {
        dfs2(u);
    }
    tout[v] = t;
}

long long sum[maxn];

long long cnt_sum[3];

long long ans = 0;

Node vals[maxn];

void add_to_root(int v, int vl, int tp) {
    vals[v].cnt[vl] += tp;
    vals[v].sum_cnt[vl] += tp;
    vals[v].sum_heavy[vl] += tp;
    int prv = v, nw = pr[v];
    while (prv != 0) {
        vals[nw].sum_cnt[vl] += tp;
        if (prv == g[nw][0]) {
            vals[nw].sum_heavy[vl] += tp;
        } else {
            if (vl != 1) {
                sum[nw] += tp * vals[prv].cnt[vl ^ 2];
            }
        }
        prv = nw;
        nw = pr[nw];
    }
}

pair<Node, long long> get_to_root(int v) {
    int prv = v, nw = pr[v];
    Node ans;
    long long ret = 0;
    while (prv != 0) {
        if (tp[nw] == 1) {
            ans = ans + vals[nw];
            ret += vals[prv].sum_cnt[tp[v] ^ 2];
        }
        prv = nw;
        nw = pr[nw];
    }
    return make_pair(ans, ret);
}

Node get_segm(int l, int r) {
    Node ans;
    for (int i = l; i < r; ++i) {
        if (tp[ord[i]] == 1) {
            ans = ans + vals[ord[i]];
        }
    }
    return ans;
}

void change(int32_t x, int32_t y) {
    if (tp[x] == y) {
        return;
    }

    ans -= cnt_all[0] * 1LL * cnt_all[1] * 1LL * cnt_all[2];
    cnt_all[tp[x]]--;
    ans += cnt_all[0] * 1LL * cnt_all[1] * 1LL * cnt_all[2];
    // cout << "? " << ans << endl;
    if (tp[x] == 0 || tp[x] == 2) {
        auto [node, ret] = get_to_root(x);
        ans += ret;
        auto node2 = get_segm(tin[0], tout[0]);
        ans += (node2.cnt[1] - node.cnt[1]) * 1LL * cnt_all[tp[x] ^ 2] - (node2.sum_cnt[tp[x] ^ 2] - node.sum_cnt[tp[x] ^ 2]);
    } else {
        // cout << "!!!" << sum[x] << ' ' << vals[x].sum_heavy[0] << ' ' << vals[x].sum_heavy[2] << endl;
        ans += sum[x];
        ans += vals[x].sum_heavy[0] * 1LL * vals[x].sum_heavy[2];
        ans += (cnt_all[0] - vals[x].sum_cnt[0]) * 1LL * (cnt_all[2] - vals[x].sum_cnt[2]);
    }
    add_to_root(x, tp[x], -1);
    // cout << "!!!" << ans << endl;

    tp[x] = y;
    ans -= cnt_all[0] * 1LL * cnt_all[1] * 1LL * cnt_all[2];
    cnt_all[tp[x]]++;
    ans += cnt_all[0] * 1LL * cnt_all[1] * 1LL * cnt_all[2];
    if (tp[x] == 0 || tp[x] == 2) {
        auto [node, ret] = get_to_root(x);
        ans -= ret;
        auto node2 = get_segm(tin[0], tout[0]);
        ans -= (node2.cnt[1] - node.cnt[1]) * 1LL * cnt_all[tp[x] ^ 2] - (node2.sum_cnt[tp[x] ^ 2] - node.sum_cnt[tp[x] ^ 2]);
    } else {
        ans -= sum[x];
        ans -= vals[x].sum_heavy[0] * 1LL * vals[x].sum_heavy[2];
        ans -= (cnt_all[0] - vals[x].sum_cnt[0]) * 1LL * (cnt_all[2] - vals[x].sum_cnt[2]);
    }
    add_to_root(x, tp[x], 1);
}

void init(int32_t n, vector<int32_t> f, vector<int32_t> u, vector<int32_t> v, int32_t q) {
    for (int i = 0; i < n; ++i) {
        tp[i] = -1;
    }
    for (int i = 0; i < n - 1; ++i) {
        g[u[i]].push_back(v[i]);
        g[v[i]].push_back(u[i]);
    }
    dfs(0);
    dfs2(0);
    ans = cnt_all[0] * 1LL * cnt_all[1] * 1LL * cnt_all[2];
    for (int x = 0; x < n; ++x) {
        tp[x] = f[x];
        ans -= cnt_all[0] * 1LL * cnt_all[1] * 1LL * cnt_all[2];
        cnt_all[tp[x]]++;
        ans += cnt_all[0] * 1LL * cnt_all[1] * 1LL * cnt_all[2];
        if (tp[x] == 0 || tp[x] == 2) {
            auto [node, ret] = get_to_root(x);
            ans -= ret;
            auto node2 = get_segm(tin[0], tout[0]);
            ans -= (node2.cnt[1] - node.cnt[1]) * 1LL * cnt_all[tp[x] ^ 2] - (node2.sum_cnt[tp[x] ^ 2] - node.sum_cnt[tp[x] ^ 2]);
            // if (x == 4) {
            //     cout << "!!!" << ret << endl;
            // }
        } else {
            ans -= sum[x];
            ans -= vals[x].sum_heavy[0] * 1LL * vals[x].sum_heavy[2];
            ans -= (cnt_all[0] - vals[x].sum_cnt[0]) * 1LL * (cnt_all[2] - vals[x].sum_cnt[2]);
        }
        add_to_root(x, tp[x], 1);
    }
}

long long num_tours() {
    // cout << "!" << ans << endl;
    return ans;
}