Colored_tree

mail@pastecode.io avatar
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