Language: C

#include <vector>
#include <stack>
#include <queue>
#include <list>
#include <bitset>
#include <set>
#include <map>
#include <unordered_set>
#include <unordered_map>
#include <algorithm>
#include <numeric>
#include <iostream>
#include <iomanip>
#include <string>
#include <chrono>
#include <random>
#include <cmath>
#include <cassert>
#include <climits>
#include <cstring>
#include <cstdlib>
#include <functional>
#include <sstream>

using namespace std;

// struct Datum {
//     public:
//         Datum() {
//         }

//         void merge(Datum& other) {
//             auto& B = other.cnts;
//             int n = cnts.size();
//             int m = B.size();
            
//             for (int i = 0, j = 0; i < n || j < m;) {

//              }
//         }

//         vector<pair<int,int>> cnts;
// };
class Solution {
public:
    vector<long long> solve(vector<int>& C, vector<pair<int,int>>& edges) {
        int n = C.size();
        vector<vector<int>> graph(n);
        for (auto& e : edges) {
            int u = e.first, v = e.second;
            graph[u].push_back(v);
            graph[v].push_back(u);
        }

        vector<int> sizes(n, 0);
        vector<long long> res(n, 0);

        vector<map<int,int>> cnts(n);

        function<void(int, int)> dfs = [&](int u, int p) {
            int w = -1;
            sizes[u] = 1;

            long long cnt = 0;
            for (auto v : graph[u]) {
                if (v == p) {
                    continue;
                }

                dfs(v, u);
                sizes[u] += sizes[v];
                cnt -= sizes[v] * 1LL * (sizes[v] - 1) / 2;
                if (w == -1 || cnts[v].size() > cnts[w].size()) {
                    w = v;
                }
            }
            cnt += sizes[u];
            cnt += (sizes[u] - 2) * 1LL * (sizes[u] - 1) / 2;
            res[C[u]] += cnt;
   
            if (w == -1) {
                cnts[u][C[u]] += 1;
                // res[C[u]] += 1;
                return;
            }

            swap(cnts[w], cnts[u]);
            
            for (auto& p : cnts[u]) {
                if (p.first != C[u]) {
                    res[p.first] += p.second * 1LL * (sizes[u] - sizes[w]); 
                }
            }
            cnts[u][C[u]] = sizes[u];

            int clr = C[u];
            for (auto v : graph[u]) {
                if (v == p || v == w) {
                    continue;
                }
                for (auto& p : cnts[v]) {
                    if (p.first != clr) {
                        long long cnt = p.second * 1LL * (sizes[u] - sizes[v]);
                        cnt -= p.second * 1LL * cnts[u][p.first];
                        res[p.first] += cnt;
                    }
                }
                for (auto& p : cnts[v]) {
                    cnts[u][p.first] += p.second;
                }
            }

            // cout << u << ":" << endl;
            // for (auto& p : cnts[u]) {
            //     cout << p.first << " " << p.second << endl;
            // } 
        };

        dfs(0, -1);
        return res;
    }
};

int main(int argc, char** argv) {
    ios::sync_with_stdio(false);
    cin.tie(0);

    int n;
    cin >> n;
    vector<int> C(n);
    for (int i = 0; i < n; ++i) {
        cin >> C[i];
        --C[i];
    }
    vector<pair<int,int>> edges;
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        cin >> u >> v;
        --u, --v;
        edges.emplace_back(u, v);
    }
    Solution sol;
    auto res = sol.solve(C, edges);
    for (int i = 0; i < n; ++i) {
        cout << res[i] << "\n";
    }
    return 0;
}

Compiler output:
main.c:1:10: fatal error: vector: No such file or directory
    1 | #include <vector>
      |          ^~~~~~~~
compilation terminated.