Language: C

#define _USE_MATH_DEFINES

#include <iostream>
#include <iomanip>
#include <stdio.h>

#include <sstream>
#include <algorithm>
#include <iterator>
#include <cmath>
#include <complex>

#include <string>
#include <cstring>
#include <vector>
#include <bitset>

#include <queue>
#include <set>
#include <map>
#include <stack>
#include <list>

#include <ctime>
////
//#include <random>//
/////////

typedef long long LL;
typedef long double LD;
typedef unsigned long long ULL;

////定数
const int INF = (int)1e9 + 10;
const LL MOD = (LL)1e9 + 7;
const LL LINF = (LL)4e18 + 20;
const LD PI = acos(-1.0);
const double EPS = 1e-9; 
/////////
using namespace::std;
template<typename T>
void chmax(T& a, T b) {
	a = max(a, b);
}

template<typename T>
void chmin(T& a, T b) {
	a = min(a, b);
}

template<typename T>
void IN(vector<T> &A,int N=-1) {
	if (N == -1) {
		N = A.size();
	}
	for (int i = 0; i < N; ++i) {
		cin >> A[i];
	}
}

template<typename T>
void SO(vector<T> &A,bool F=false) {
	sort(A.begin(), A.end());
	if (F) {
		reverse(A.begin(), A.end());
	}
}
//////////////
class BITree {//1-index
	int N;
	vector<LL> bit;
public:
	BITree(int n) {
		N = n;
		bit = vector<LL>(N + 1, 0);//1-index
	}
	void add(int a, LL w) {//aにwを足す
		if (a <= 0 || N < a) return;//a:[1,N]
		for (int i = a; i <= N; i += i & -i) {
			bit[i] += w;
		}
	}
	LL sum(int a) {//[1,a]の和,a:[1,N]
		/*
		1番目からa番目までの和、1-index
		*/
		LL ret = 0;
		if (a > N) a = N;
		for (int i = a; i > 0; i -= i & -i) {
			ret += bit[i];
		}
		return ret;
	}
};

vector<vector<int>> G;
int k = 0;//
vector<int> in, out;

void dfs(int u, int P) {
	in[u] = k++;
	for (int v : G[u]) {
		if (v == P)continue;
		dfs(v, u);
	}
	out[u] = k;
}

LL c2(LL n) {
	return n * (n + 1) / 2;
}
void solve() {
	int N;
	cin >> N;

	vector<int> C(N);
	G = vector<vector<int>>(N);
	in = vector<int>(N);
	out = vector<int>(N);

	for (int i = 0; i < N; ++i) {
		cin >> C[i];
		C[i]--;
	}
	for (int i = 0; i < N - 1; ++i) {
		int a, b;
		cin >> a >> b;
		a--; b--;
		G[a].push_back(b);
		G[b].push_back(a);
	}

	dfs(0, -1);
	vector<vector<int>> vs(N);//vs[c]:色cの頂点集合
	for (int i = 0;i<N;++i){
		vs[C[i]].push_back(i);
	}

	//色cの頂点集合を、行きがけ順が遅い順にソート。
	for (int c = 0; c < N; ++c) {
		sort(vs[c].begin(), vs[c].end(),
			[&](int a, int b) {return in[a] > in[b]; }
			);
	}

	BITree bit(N);//1-index
	for (int i = 1; i <= N; ++i) {
		bit.add(i, 1);
	}
	
	for (int c = 0; c < N; ++c) {
		//vs[c]
		LL ans = c2(N);
		vector<pair<int, int> > his;//bitの操作情報

		for (int v : vs[c]) {
			int cnt = 1;
			for (int u : G[v]) {
				//vより速い物は親なのでスキップ
				if (in[u] < in[v])continue;

				//outは次の番号になっているので
				//[in[u],out[u])の区間
				/*
				inにはL,outにはR+1が入っている。
				[L,R+1)->[L+1,R+2)->[R+2-1]-[L+1-1]:=[R+1]-[L]
				*/
				int num = bit.sum(out[u]) - bit.sum(in[u]);
				ans -= c2(num);
				cnt += num;//
			}
			bit.add(in[v] + 1, -cnt);
			his.emplace_back(in[v] + 1, cnt);
		}
		ans -= c2(bit.sum(N));//rootの分
		//rootが指定頂点なら0
		//指定色でないなら残った頂点数。
		for (auto p : his) {
			bit.add(p.first, p.second);
		}
		cout << ans << endl;
	}
}

signed main(void) {
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);
	std::cout << std::fixed;//小数を10進数表示
	cout << setprecision(16);//小数点以下の桁数を指定//coutとcerrで別	

	solve();
}


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