Language: C++

#include <bits/stdc++.h>
 
#define _overload(_1,_2,_3,name,...) name
#define _rep(i,n) _range(i,0,n)
#define _range(i,a,b) for(int i=(int)(a);i<(int)(b);++i)
#define rep(...) _overload(__VA_ARGS__,_range,_rep,)(__VA_ARGS__)
 
#define _rrep(i,n) _rrange(i,0,n)
#define _rrange(i,a,b) for(int i=(int)(b)-1;i>=(int)(a);--i)
#define rrep(...) _overload(__VA_ARGS__,_rrange,_rrep,)(__VA_ARGS__)
 
#define _all(arg) begin(arg),end(arg)
#define uniq(arg) sort(_all(arg)),(arg).erase(unique(_all(arg)),end(arg))
#define bit(n) (1LL<<(n))
 
#define dbg(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ") " << __FILE__ << endl;

using namespace std;

// overload: >>, <<
template<typename T> istream &operator>>(istream &is, vector<T> &vec){ for (auto &v : vec) is >> v; return is; }
template<typename T> ostream &operator<<(ostream &os, const vector<T> &vec){ os << "["; for (auto v : vec) os << v << ","; os << "]"; return os; }
template<typename T> ostream &operator<<(ostream &os, const deque<T> &vec){ os << "deq["; for (auto v : vec) os << v << ","; os << "]"; return os; }
template<typename T> ostream &operator<<(ostream &os, const set<T> &vec){ os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template<typename T> ostream &operator<<(ostream &os, const unordered_set<T> &vec){ os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template<typename T> ostream &operator<<(ostream &os, const multiset<T> &vec){ os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template<typename T> ostream &operator<<(ostream &os, const unordered_multiset<T> &vec){ os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template<typename T1, typename T2> ostream &operator<<(ostream &os, const pair<T1, T2> &pa){ os << "(" << pa.first << "," << pa.second << ")"; return os; }
template<typename TK, typename TV> ostream &operator<<(ostream &os, const map<TK, TV> &mp){ os << "{"; for (auto v : mp) os << v.first << "=>" << v.second << ","; os << "}"; return os; }
template<typename TK, typename TV> ostream &operator<<(ostream &os, const unordered_map<TK, TV> &mp){ os << "{"; for (auto v : mp) os << v.first << "=>" << v.second << ","; os << "}"; return os; }

// overload: +, -, *, /
template<typename T1, typename T2> pair<T1, T2> operator+(const pair<T1, T2> &l, const pair<T1, T2> &r) { return make_pair(l.first + r.first, l.second + r.second); }
template<typename T1, typename T2> pair<T1, T2> operator-(const pair<T1, T2> &l, const pair<T1, T2> &r) { return make_pair(l.first - r.first, l.second - r.second); }

// template functions
template<typename T> void ndarray(vector<T> &vec, int len) { vec.resize(len); }
template<typename T, typename... Args> void ndarray(vector<T> &vec, int len, Args... args) { vec.resize(len); for (auto &v : vec) ndarray(v, args...); }
template<class T>bool chmax(T &a, const T &b) { return (a<b)?(a=b,1):0;}
template<class T>bool chmin(T &a, const T &b) { return (b<a)?(a=b,1):0;}
 
using ll=long long;
using vi=vector<int>;
using vll=vector<ll>;
 
const double EPS = 1e-10;
const double PI = acos(-1.0);
const int inf = INT_MAX;
const ll INF = LLONG_MAX;
const ll MOD = 1000000007LL;
const int dx[4]={1,0,-1,0};
const int dy[4]={0,1,0,-1};
 
ll extgcd(ll a,ll b,ll& x,ll& y){x=1,y=0;ll g=a;if(b!=0) g=extgcd(b,a%b,y,x),y-=a/b*x;return g;}
ll ADD(const ll &a, const ll &b, const ll mod=MOD) { return (a+b)%mod;}
ll SUB(const ll &a, const ll &b, const ll mod=MOD) { return (a-b+mod)%mod;}
ll MUL(const ll &a, const ll &b, const ll mod=MOD) { return (1LL*a*b)%mod;}
ll DIV(const ll &a, const ll &b, const ll mod=MOD) {ll x,y; extgcd(b,mod,x,y);return MUL(a,(x+mod)%mod,mod);}
 
int n; 
vi C;
vector<vi> G;
vi subt;

vi pa_v;
vi pa_ci;

vi c2v;
vi c2ci;

int dfs_subt(int v, int pv){
    int& cur = subt[v];
    cur = 1;
    for(auto& nv : G[v]){
        if(nv == pv) continue;
        cur += dfs(nv, v);
    }
    return cur;
}

void dfs_pa(int v, int pv){
    int c = C[v];
    int par_ci = c2ci[c], par_v = c2v[c];
    if(par_v != -1){
        pa_v[v] = par_v;
        pa_ci[v] = par_ci;
    }

    c2v[c] = v;
    rep(i, G[v].size()){
        int nv = G[v][i];
        if(nv == pv) continue;
        c2ci[c] = i;
        dfs_pa(nv, v);
    }
}

vll dp;
vector<vi> used_cnt;
vi nevuse_cnt;
void dfs(int v, int pv){
    int c = C[v];

    used_cnt[v] = vi(G[v].size());

    // v --> v
    ll& cur = dp[c];
    cur += 1;

    rep(i, G[v].size()){
        int nv = G[v][i];
        if(nv == pv) continue;
        dfs(nv, v);
    }
    int sum = subt[v] - 1;
    rep(i, G[v].size()){
        sum -= used_cnt[v][i];
    }

    // subt <--> subt
    rep(i, G[v].size()){
        int nv = G[v][i];
        ll ccnt = subt[nv] - used_cnt[v][i];
        cur += 2LL * ccnt * (sum - ccnt);
    }

    // subt <--> parent_t
    rep(i, G[v].size()){
        int nv = G[v][i];
        ll ccnt = subt[nv] - used_cnt[v][i];
        ll pcnt = n - nevuse_cnt[c] - ccnt;
        cur += 2LL * ccnt * pcnt;
    }

    nevuse_cnt[c] += sum + 1;
    int par_v = pa_v[v];
    int par_ci = pa_ci[v];
    used_cnt[par_v][par_ci] += sum + 1;
}

int main(void){
    cin.tie(0);
    ios::sync_with_stdio(false);

    cin >> n;
    C = vi(n); for(auto& e : C) cin >> e, e--;

    G.resize(n);
    rep(_, n - 1){
        int a, b; cin >> a >> b;
        a--, b--;
        G[a].emplace_back(b);
        G[b].emplace_back(a);
    }

    subt = vi(n);
    dfs_subt(0, -1);

    pa_ci = pa_v = c2v = c2ci = vi(n, -1);
    dfs_pa(0, -1);

    dp = vll(n);
    used_cnt.resize(n);
    dfs(0, -1);

    return 0;
}


Compiler output:
a.cc: In function 'int dfs_subt(int, int)':
a.cc:76:16: error: 'dfs' was not declared in this scope; did you mean 'ffs'?
   76 |         cur += dfs(nv, v);
      |                ^~~
      |                ffs