Submission #35203807

---

Source Code Expand

#include <cassert>
#include <iostream>
#include <vector>

struct graph : std::vector<std::vector<int>> {
    using std::vector<std::vector<int>>::vector;
    void add_edge(int u, int v, bool directed = false) {
        (*this)[u].emplace_back(v);
        if (directed) return;
        (*this)[v].emplace_back(u);
    }
};

int main() {
    int n;
    std::cin >> n;
    graph g(n);
    std::vector<int> a(n);
    for (auto &val : a) std::cin >> val;
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        std::cin >> u >> v;
        u--;
        v--;
        g.add_edge(u, v);
    }
    std::vector dp(n, std::vector<int>(2, 0));
    auto dfs = [&](auto &&self, int v, int par = -1) -> void {
        if (par != -1 && g[v].size() == 1) {
            dp[v][1 ^ a[v]] = 1;
            return;
        }
        bool flag = false;
        for (auto nv : g[v]) {
            if (nv == par) continue;
            self(self, nv, v);
            auto nxt = dp[v];
            if (!flag) {
                dp[v] = dp[nv];
                flag = true;
                continue;
            }
            nxt[0] = std::max(dp[v][0] + dp[nv][0], dp[v][1] + dp[nv][1]);
            nxt[1] = std::max(dp[v][0] + dp[nv][1], dp[v][1] + dp[nv][0]);
            dp[v] = nxt;
        }
        dp[v][1 ^ a[v]]++;
        return;
    };
    dfs(dfs, 0);
    std::cout << std::max(dp[0][0], dp[0][1]) << '\n';
}

Submission Info

Judge Result