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;

template <typename T>
struct EdgeT {
    EdgeT(int _u = -1, int _v = -1, T _w = 0) : u(_u), v(_v), w(_w) {
    }

    int u, v;
    T w;
};

using Edge = EdgeT<int>;


class Solution {
public:
    vector<long long> solve(vector<int>& C, vector<pair<int,int>>& E) {
        int n = C.size();
        init(n);

        for (auto& e : E) {
            int u = e.first;
            int v = e.second;
            graph[u].push_back(edges.size());
            graph[v].push_back(edges.size());
            edges.emplace_back(u, v, 0LL);
        }

        rt = 0;
        clock = 0;
        dfs(rt, -1, clock, 0, 0);

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

        vector<vector<int>> que(n);
        vector<int> cnts(n, 0);
        function<void(int,int)> calc = [&](int u, int p) {
            int clr = C[u];
            auto& Q = que[clr];

            int sum = 0;
            for (auto e : graph[u]) {
                int v = edges[e].u + edges[e].v - u;
                if (v == p) {
                    continue;
                }
                calc(v, u);
                int cnt = size[v];
                
                int l = 0, r = Q.size();
                while (l < r) {
                    int m = (l + r) >> 1;
                    if (in[Q[m]] < in[v]) {
                        l = m + 1;
                    } else {
                        r = m;
                    }
                }
                for (int i = l; i < Q.size(); ++i) {
                    if (isAncestor(v, Q[i])) {
                        cnt -= size[Q[i]];
                    } else {
                        break;
                    }
                }
                res[clr] += cnt * 1LL * sum;
                sum += cnt;
            }



            int down = size[u];
            int up = n - size[u] + 1;
            while (!Q.empty()) {
                int v = Q.back();
                // cout << clr << " " << v << " " << u << endl;
                if (isAncestor(u, v)) {
                    down -= size[v];
                    cnts[clr] -= size[v];
                    Q.pop_back();
                } else {
                    break;
                }
            }
            Q.push_back(u);
            up -= cnts[clr];
            cnts[clr] += size[u];
            // cout << u << " " << up << " " << down << endl;
            res[clr] += down * 1LL * up;
            
        };

        calc(rt, -1);


        return res;
    }

private:
    int N;
    int rt;
    int clock;
    vector<Edge> edges;
    vector<vector<int>> graph;
    vector<vector<int>> pars;
    vector<int> depth;
    vector<long long> dis;
    vector<int> size;
    vector<int> in;
    vector<int> out;

    void init(int n) {
        N = n;
        rt = -1;
        clock = 0;
        edges.clear();
        graph.clear();
        graph.resize(n);
        depth.clear();
        depth.resize(n, 0);
        dis.clear();
        dis.resize(n, 0);
        size.clear();
        size.resize(n, 0);
        in.clear();
        in.resize(n, -1);
        out.clear();
        out.resize(n, -1);

        {
            int h = 0;
            while (1 << (h + 1) <= n) {
                ++h;
            }
            pars.clear();
            pars.resize(n, vector<int>(h + 1, -1));
        }
    }
    
    void dfs(int u, int p, int& clock, int dep, long long d) {
        pars[u][0] = p;
        {
            int h = pars[u].size();
            for (int i = 1; i < h; ++i) {
                if (pars[u][i - 1] == -1) {
                    break;
                }
                pars[u][i] = pars[pars[u][i - 1]][i - 1];
            }
        }
        depth[u] = dep;
        dis[u] = d;
        in[u] = ++clock;

        size[u] = 1;
        for (auto e : graph[u]) {
            int v = edges[e].u + edges[e].v - u;
            auto w = edges[e].w;
            if (v == p) {
                continue;
            }
            dfs(v, u, clock, dep + 1, d + w);
            size[u] += size[v];
        }

        out[u] = ++clock;
    }

    int lca(int u, int v) {
        if (u == -1 || v == -1) {
            return u == -1 ? v : u;
        }
        if (u == v) {
            return u;
        }
        if (depth[u] > depth[v]) {
            swap(u, v);
        }
        int h = pars[v].size();
        for (int i = h - 1; i >= 0; --i) {
            int p = pars[v][i];
            if (p != -1 && depth[p] >= depth[u]) {
                v = p;
            }
        }
        if (u == v) {
            return u;
        }
        for (int i = h - 1; i >= 0; --i) {
            if (pars[u][i] != pars[v][i]) {
                u = pars[u][i];
                v = pars[v][i];
            }
        }
        return pars[u][0];
    }

    bool isAncestor(int u, int v) {
        return in[u] <= in[v] && out[v] <= out[u];
    }

};
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.