Submission #54843426

---

Source Code Expand

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define ld long double
#define fi first
#define se second
#define pb push_back
#pragma GCC optimize("unroll-loops,O3")
#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#define cok cout << (ok ? "YES\n" : "NO\n");
#define dbg(x) cout << (#x) << ": " << x << endl;
#define dbga(x,l,r) cout << (#x) << ": "; for (int ii=l;ii<r;ii++) cout << x[ii] << " "; cout << endl;
#define int long long
#define pi pair<int, int>
const int N = 1e6+9;
int n, ans, a[N], sz[N];
vector<int> vec[N];
map<int, pi> *mp[N];
int add[N];
void dfs(int v, int p = -1) {
    sz[v] = 1;
    pi mx = {-1, -1};
    for (int u: vec[v]) {
        if (u == p) continue;
        dfs(u, v);
        sz[v] += sz[u];
        mx = max(mx, pi{sz[u], u});
    }
    if (mx.se == -1) {
        mp[v] = new map<int, pi>;
        (*mp[v])[a[v]] = {1, 0};
        return;
    }
    mp[v] = mp[mx.se];
    add[v] = add[mx.se] + 1;
    ans += (*mp[v])[a[v]].se + (*mp[v])[a[v]].fi * add[v];
    // dbg(a[v]);
    // dbg((*mp[v])[a[v]].se + (*mp[v])[a[v]].fi * add[v]);
    (*mp[v])[a[v]].fi++;
    (*mp[v])[a[v]].se -= add[v];
    for (int u: vec[v]) {
        if (u == mx.se || u == p) continue;
        for (auto [color, para]: *mp[u]) {
            // dbg(color << " " << v + 1);
            // dbg(((*mp[v])[color].se + (*mp[v])[color].fi * add[v]) * (para.se + para.fi * (add[u] + 1)));
            ans += ((*mp[v])[color].se + (*mp[v])[color].fi * add[v]) * para.fi + (para.se + para.fi * (add[u] + 1)) * (*mp[v])[a[v]].fi;
        }
        for (auto [color, para]: *mp[u]) {
            (*mp[v])[color].fi += para.fi;
            (*mp[v])[color].se += para.se + para.fi * add[u] - para.fi * add[v];
        }
    }
    // dbg(v + 1);
    // dbg(ans);
}
signed main() {
    cin.tie(0); ios_base::sync_with_stdio(0);
    cin >> n;
    for (int i = 1; i < n; i++) {
        int a, b;
        cin >> a >> b;
        a--;b--;
        vec[a].pb(b);
        vec[b].pb(a);
    }
    for (int i = 0; i < n; i++) cin >> a[i];
    dfs(0);
    cout << ans;
}

Submission Info

Judge Result