#include "joitour.h"
#include <bits/stdc++.h>

using namespace std;
using ll = long long;

constexpr int N = 2e5 + 2;

int n;
int type[N];
vector<int> adj[N];

namespace Tree {
    struct Fenwick {
        int fn[N]{};

        void modify(int x, int t) {
            for (++x; x < N; x += x & -x) {
                fn[x] += t;
            }
        }

        void modify(int l, int r, int t) {
            modify(l, t);
            modify(r, -t);
        }

        int sum(int x) {
            int sum = 0;
            for (++x; x > 0; x -= x & -x) {
                sum += fn[x];
            }
            return sum;
        }

        int sum(int l, int r) {
            return sum(r - 1) - sum(l - 1);
        }
    };

    int tin[N], tout[N], jump[N], depth[N], father[N];
    Fenwick fn[3], up;

    bool isp(int a, int b) {
        return tin[a] <= tin[b] && tout[a] >= tout[b];
    }

    void modify(int x, int val, int coef) {
//        cout << "modify: " << x << " " << val << " " << coef << endl;
        if (val == 1) {
            up.modify(tin[x], tout[x], coef);
        }
        fn[val].modify(tin[x], coef);
    }

    void init() {
        int timer = 0;
        auto dfs = [&](auto self, int v) -> void {
            tin[v] = timer++;
            for (int to : adj[v]) {
                if (to != father[v]) {
                    father[to] = v;
                    depth[to] = depth[v] + 1;
                    int p = jump[v];
                    int pp = jump[p];
                    if (depth[pp] - depth[p] == depth[p] - depth[v]) {
                        jump[to] = pp;
                    } else {
                        jump[to] = v;
                    }
                    self(self, to);
                }
            }
            tout[v] = timer;
        };
        dfs(dfs, 0);
    }

    int lca(int a, int b) {
        while (!isp(a, b)) {
            if (!isp(jump[a], b)) {
                a = jump[a];
            } else {
                a = father[a];
            }
        }
        return a;
    }

    int orientedSum(int x, int f, int s) {
        if (!isp(x, f)) {
            return fn[s].sum(tin[x], tout[x]);
        } else {
            while (!isp(father[f], x)) {
                if (!isp(father[jump[f]], x)) {
                    f = jump[f];
                } else {
                    f = father[f];
                }
            }
            return fn[s].sum(0, n) - fn[s].sum(tin[f], tout[f]);
        }
    }

    int pathOm(int a) {
        return up.sum(tin[a]);
    }

    int countOm(int a, int b) {
        int c = lca(a, b);
        return pathOm(a) + pathOm(b) - pathOm(c) * 2 + (fn[1].sum(tin[c], tin[c] + 1));
    }
}

vector<array<ll, 4>> kids[N];
array<ll, 4> sums[N];
ll sumProd[N];
vector<pair<int, int>> parents[N]; // vertex, pos of my dad in kids[vertex]
bool used[N];
int siz[N];

void findSz(int v, int par) {
    siz[v] = 1;
    for (int to : adj[v]) {
        if (to != par && !used[to]) {
            findSz(to, v);
            siz[v] += siz[to];
        }
    }
}

int centroid(int x, int m, int par) {
    for (int to : adj[x]) {
        if (to != par && !used[to] && siz[to] * 2 > m) {
            return centroid(to, m, x);
        }
    }
    return x;
}

void build(int x) {
    findSz(x, -1);
    auto dfs1 = [&](auto self, int v, int par) -> void {
        parents[v].push_back({x, kids[x].size() - 1});
        for (int to : adj[v]) {
            if (to != par && !used[to]) {
                self(self, to, v);
            }
        }
    };
    for (int to : adj[x]) {
        if (!used[to]) {
            kids[x].push_back({});
            dfs1(dfs1, to, x);
        }
    }
    used[x] = true;
    for (int to : adj[x]) {
        if (!used[to]) {
            build(centroid(to, siz[to], -1));
        }
    }
}

int toid(int x) {
    return x == -1 ? -2 : x == 0 ? 0 : x == 1 ? -1 : 1;
}

ll countTours(int v, int coef) {
    if (v == 4) {
        cout << "here!" << endl;
    }
    Tree::modify(v, type[v], coef);
    if (type[v] == 1) {
        ll ans = 0;
        for (auto [x, p] : parents[v]) {
            for (int id = 0; id < 2; ++id) {
                int s = id * 2;
                int sub = Tree::orientedSum(v, x, s);
                int oth = sums[x][2 + (id ^ 1)] - kids[x][p][(2 + (id ^ 1))] + (toid(type[x]) == (id ^ 1));
                ans += 1LL * sub * oth;
                sums[x][id] += sub * coef;
                kids[x][p][id] += sub * coef;
            }
        }
        // when I'm the root
        ans += 1LL * sums[v][2] * sums[v][3];
        ans -= sumProd[v];
        return ans * coef;
    }
    int id = type[v] == 0 ? 0 : 1;
    ll ans = 0;
    for (auto [x, p] : parents[v]) {
        sumProd[x] -= kids[x][p][2] * kids[x][p][3];
        ans += sums[x][id ^ 1] - kids[x][p][id ^ 1]; // when the mid is not in our subtree
        int mid = Tree::countOm(x, v);
        ans += 1LL * mid * (sums[x][(id ^ 1) + 2] - kids[x][p][(id ^ 1) + 2] + (toid(type[x]) == (id ^ 1)));
        int now = mid - (type[x] == 1); // initially put -1 everywhere
        kids[x][p][id] += now * coef;
        sums[x][id] += now * coef;
        kids[x][p][2 + id] += coef;
        sums[x][2 + id] += coef;
        sumProd[x] += kids[x][p][2] * kids[x][p][3];
    }
    // when I'm the root
    ans += sums[v][id ^ 1];
    return ans * coef;
}

ll answer = 0;

void init(int N_, std::vector<int> F, std::vector<int> U, std::vector<int> V,
          int Q) {
    n = N_;
    fill(type, type + n, -1);
    for (int i = 0; i < n - 1; ++i) {
        adj[U[i]].push_back(V[i]);
        adj[V[i]].push_back(U[i]);
    }
    Tree::init();
    findSz(0, -1);
    build(centroid(0, n, -1));
    for (int i = 0; i < n; ++i) {
        type[i] = F[i];
        ll now = countTours(i, 1);
//        cout << now << endl;
        answer += now;
    }
}

void change(int X, int Y) {
    answer += countTours(X, -1);
    type[X] = Y;
    answer += countTours(X, 1);
}

long long num_tours() {
    return answer;
}