Colored_tree
unknown
c_cpp
a year ago
1.8 kB
11
Indexable
#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';
}
}
Editor is loading...
Leave a Comment