OIG X - jablko

// https://szkopul.edu.pl/problemset/problem/ZXKiCopRuLe5diKPU0oMNgNn/site/?key=statement

#include <bits/stdc++.h>

using namespace std;

#define int long long

constexpr int sizik = 200 * 1001;

#define ar std::array
#define pr std::pair
#define vec std::vector

typedef vec<vec<int>> _kra;

int n;
int sum = 0, wystuplenia[sizik];

void remove(int idx);
void add(int idx);
int get_answer();

int block_size = 317;

struct Query {
    int l, r, idx;
    bool operator<(Query other) const { return make_pair(l / block_size, r) < make_pair(other.l / block_size, other.r); }
    Query(int l1, int r1, int idx1) : l(l1), r(r1), idx(idx1) { ; }
};

int answers[sizik];
vector<Query> queries;
void mo_s_algorithm() {
    sort(queries.begin(), queries.end());

    sum = 0;
    for (int i = 0; i <= n; i++) {
        wystuplenia[i] = 0;
    }

    int cur_l = 1;
    int cur_r = 0;
    // invariant: data structure will always reflect the range [cur_l, cur_r]
    for (Query q : queries) {
        while (cur_l > q.l) {
            cur_l--;
            add(cur_l);
        }
        while (cur_r < q.r) {
            cur_r++;
            add(cur_r);
        }
        while (cur_l < q.l) {
            remove(cur_l);
            cur_l++;
        }
        while (cur_r > q.r) {
            remove(cur_r);
            cur_r--;
        }
        answers[q.idx] = get_answer();
    }
}

int pre[sizik], post[sizik], k[sizik], k1[sizik];
std::vector<int> kra[sizik];
int counter = 1;
void DFS(int v, int p) {
    pre[v] = counter++;

    for (const auto& u : kra[v]) {
        if (u != p) {
            DFS(u, v);
        }
    }

    post[v] = counter++;
}

void solve() {
    std::cin >> n;

    for (int i = 1; i <= n; i++) {
        std::cin >> k[i];
    }

    for (int i = 0; i < n - 1; i++) {
        int a, b;
        std::cin >> a >> b;

        kra[a].push_back(b);
        kra[b].push_back(a);
    }

    DFS(1, 1);
    for (int i = 1; i <= n; i++) {
        kra[i].clear();
    }

    for (int i = 1; i <= n; i++) {
        k1[pre[i]] = k[i];
        k1[post[i]] = k[i];
    }

    for (int i = 1; i <= n; i++) {
        queries.push_back({pre[i], post[i], i});
    }

    mo_s_algorithm();

    for (int i = 1; i <= n; i++) {
        std::cout << answers[i] << " ";
    }
    std::cout << '\n';
}

void remove(int idx) {
    sum -= 2 * (--wystuplenia[k1[idx]]) + 1;
}

void add(int idx) {
    sum += 2 * (wystuplenia[k1[idx]]++) + 1;
}

int get_answer() {
    return sum >> 2;
}

int32_t main() {
    std::ios_base::sync_with_stdio(0);
    std::cin.tie(0);
    std::cout.tie(0);

    int t = 1;
    // std::cin >> t;

    for (; t > 0; t--) {
        solve();
    }

    return 0;
}