Colored_tree
unknown
c_cpp
a month ago
1.8 kB
3
Indexable
Never
#include <bits/stdc++.h> using namespace std; typedef long long ll; #define f first #define s second const int N = 2e5; int color[N] , dep[N] , par[N]; vector<int> adj[N]; void dfs(int u = 0 , int p = -1){ par[u] = p; for(auto &v : adj[u]) if(v != p){ dep[v] = dep[u] + 1; dfs(v , u); } } int sz , path[N]; void genPath(int u , int v){ if(dep[u] < dep[v]) swap(u , v); sz = 0; while(dep[u] != dep[v]){ path[sz++] = u; u = par[u]; } while(u != v){ path[sz++] = u; path[sz++] = v; u = par[u]; v = par[v]; } path[sz++] = u; } ll totFrq , totC , totC2 , ans; struct info{ ll frq , sumC , sumC2; void apply(int c , int f){ ll cc = c * c; // f = 1 -> insert , f = -1 -> delete ans += f * (totFrq - frq) * cc; ans += f * (totC2 - sumC2); ans += -f * 2LL * (totC - sumC) * c; frq += f; sumC += f * c; sumC2 += f * cc; totFrq += f; totC += f * c; totC2 += f * cc; } }I[N]; void add(int c){ I[c].apply(c , 1); } void erase(int c){ I[c].apply(c , -1); } ll solve(int u , int v){ genPath(u , v); for(int i = 0; i < sz; ++i) add(color[path[i]]); ll cnt = ans; for(int i = 0; i < sz; ++i) erase(color[path[i]]); return cnt; } int main(){ ios::sync_with_stdio(0); cin.tie(NULL); cout.tie(0); int n , q; cin >> n >> q; for(int i = 0; i < n; ++i) cin >> color[i]; for(int i = 1; i < n; ++i){ int u, v; cin >> u >> v; --u; --v; adj[u].emplace_back(v); adj[v].emplace_back(u); } dfs(); while(q--){ int u, v; cin >> u >> v; --u; --v; cout << solve(u , v) << '\n'; } }
Leave a Comment