Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code: Closer

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

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int T;
    cin >> T;
    while (T--) {
        int n;
        cin >> n;
        vector<int> wt(n), a(2 * n);
        for (auto &x : wt)
            cin >> x;
        for (auto &x : a)
            cin >> x, x--;
        vector<vector<int>> G(2 * n);
        vector<int> partner(2 * n), lst(n, -1);
        for (int i = 0; i < 2 * n; i++) {
            if (lst[a[i]] != -1) {
                partner[i] = lst[a[i]];
                partner[lst[a[i]]] = i;
            }
            lst[a[i]] = i;
        }

        for (int i = 0; i < 2 * n - 1; i++) {
            int u, v;
            cin >> u >> v;
            u--, v--;
            G[u].push_back(v);
            G[v].push_back(u);
        }
        vector<int64_t> par_dp(2 * n); // v -> best if swap(v, par[v]) is made
        vector take_dp(2 * n,
                       vector<pair<int, int64_t>>()); // v -> list of (val at v,
                                                      // best) if swap(v, child)
                                                      // is made for some child
        vector<int64_t> stay_dp(2 * n); // v -> best if no swap is made at v

        vector<int64_t> default_best(2 * n);
        vector<int> par(2 * n, -1);
        auto dfs = [&](this auto dfs, int v, int p) -> void {
            if (p != -1)
                G[v].erase(find(G[v].begin(), G[v].end(), p));
            par[v] = p;
            int64_t s = 0;
            for (auto &u : G[v]) {
                dfs(u, v);
                int64_t best_stay = stay_dp[u] + (a[u] == a[v] ? wt[a[v]] : 0);
                int64_t best_par =
                    stay_dp[u] + (p != -1 && a[p] == a[u] ? wt[a[p]] : 0);
                default_best[u] = stay_dp[u];
                for (auto &[x, val] : take_dp[u]) {
                    best_stay =
                        max(best_stay, val + (x == a[v] ? wt[a[v]] : 0));
                    best_par = max(best_par,
                                   val + (p != -1 && x == a[p] ? wt[a[p]] : 0));
                    default_best[u] = max(default_best[u], val);
                }
                stay_dp[v] += best_stay;
                par_dp[v] += best_par;
                s += default_best[u];
            }

            for (auto &u : G[v]) {
                int64_t score =
                    s - default_best[u] +
                    par_dp[u]; // if a[u] == a[v] there's no point in swapping
                // find a match
                int w = partner[u];
                if (par[w] != -1) {
                    if (par[w] == v) {
                        score +=
                            max(0LL, stay_dp[w] + wt[a[u]] - default_best[w]);
                    } else if (par[w] != u && par[par[w]] == v) {
                        auto it = lower_bound(take_dp[par[w]].begin(),
                                              take_dp[par[w]].end(),
                                              make_pair(a[w], -1LL));
                        assert(it != take_dp[par[w]].end() &&
                               it->first == a[w]);
                        score += max(0LL, it->second + wt[a[u]] -
                                              default_best[par[w]]);
                    }
                }
                take_dp[v].emplace_back(a[u], score);
            }
            sort(take_dp[v].begin(), take_dp[v].end());
        };
        dfs(0, -1);
        int64_t ans = stay_dp[0];
        for (auto &[x, val] : take_dp[0])
            ans = max(ans, val);
        cout << ans << '\n';
    }
}