1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
| #include <bits/stdc++.h> using namespace std; typedef long long LL; #define FOR(i, x, y) for (decay<decltype(y)>::type i = (x), _##i = (y); i < _##i; ++i) #define FORD(i, x, y) for (decay<decltype(x)>::type i = (x), _##i = (y); i > _##i; --i) #ifdef zerol #define dbg(args...) do { cout << "\033[32;1m" << #args<< " -> "; err(args); } while (0) #else #define dbg(...) #endif void err() { cout << "\033[39;0m" << endl; } template<typename T, typename... Args> void err(T a, Args... args) { cout << a << ' '; err(args...); }
const int maxn = 1E5 + 100, maxc = 1E5 + 100; int c[maxn], in[maxn], out[maxn], sz[maxn], pa[maxn]; vector<int> G[maxn], C[maxn]; int n, clk = 0;
struct BIT { LL c[maxn]; inline int lowbit(int x) { return x & -x; } void add(int x, int k) { for (int i = x; i <= n; i += lowbit(i)) c[i] += k; } LL sum(int x) { LL ret = 0; for (int i = x; i > 0; i -= lowbit(i)) ret += c[i]; return ret; } void add(int l, int r, int v) { add(l, v); add(r + 1, -v); } } bit;
void init_dfs(int u, int fa) { in[u] = clk++; C[c[u]].push_back(u); pa[u] = fa; sz[u] = 1; for (int v: G[u]) { if (v == fa) continue; init_dfs(v, u); sz[u] += sz[v]; } out[u] = clk - 1; }
void go(const vector<int>& V, int& k) { vector<int> nxt; int u = V[k]; for (int v: G[u]) { if (v == pa[u]) continue; int num = sz[v]; nxt.clear(); while (k + 1 < V.size()) { int to = V[k + 1]; if (in[to] <= out[v]) { nxt.push_back(to); num -= sz[to]; go(V, ++k); } else break; } bit.add(in[v], out[v], num); for (int to: nxt) bit.add(in[to], out[to], -num); } }
int main() { #ifdef zerol freopen("in", "r", stdin); #endif cin >> n;
FOR (i, 1, maxc) C[i].push_back(0); G[0].push_back(1); pa[1] = 0; c[0] = 0;
FOR (i, 1, n + 1) scanf("%d", &c[i]); FOR (_, 1, n) { static int u, v; scanf("%d%d", &u, &v); G[u].push_back(v); G[v].push_back(u); } init_dfs(0, -1);
int tmp; LL color_cnt = 0; FOR (k, 1, maxc) { if (C[k].size() == 1) continue; color_cnt++; go(C[k], tmp = 0); } FOR (i, 1, n + 1) printf("%lld\n", color_cnt * n - bit.sum(in[i])); }
|