Language: C

#include <iostream>
#include <string>
#include <algorithm>
#include <vector>
#include <cmath>
#include <queue>
#include <map>
#include <unordered_map>
#include <set>
#include <functional>
#include <bitset>
#include <numeric>
#include <complex>
#include <cassert>
#include <random>


class SegmentTree {
private:
	using Monoid = int;
	const Monoid INI = 0;
	Monoid op(Monoid a, Monoid b) {
		return a + b;
	}

	int num;
	std::vector<Monoid> dat;

public:
	SegmentTree(int n_) {
		num = 1;
		while (num < n_) num *= 2;
		dat = std::vector<Monoid>(2 * num, INI);
	}

	SegmentTree(const std::vector<Monoid> &m) : SegmentTree(m.size()) {
		int n = m.size();
		for (int i = 0; i < n; i++) {
			dat[num - 1 + i] = m[i];
		}
		for (int i = num - 2; i >= 0; i--) {
			dat[i] = op(dat[2 * i + 1], dat[2 * i + 2]);
		}
	}

	// Update the k-th value (0-indexed) to a.
	void update(int k, Monoid a) {
		k += num - 1;
		dat[k] = a;
		while (k > 0) {
			k = (k - 1) / 2;
			dat[k] = op(dat[2 * k + 1], dat[2 * k + 2]);
		}
	}

	// Get the value of the range [a, b).
	// k: The index of the node.
	// [left, right): The range corresponds to the k-th node.
	// Call like getval(a, b).
	Monoid getval(int a, int b, int k = 0, int left = 0, int right = -1) {
		if (right < 0) right = num;
		if (right <= a || b <= left) return INI;
		if (a <= left && right <= b) return dat[k];
		Monoid vleft = getval(a, b, 2 * k + 1, left, left + (right - left) / 2);
		Monoid vright = getval(a, b, 2 * k + 2, left + (right - left) / 2, right);
		return op(vleft, vright);
	}
};



int main() {
	int n; std::cin >> n;
	std::vector<int> color(n);
	for (int i = 0; i < n; i++) {
		std::cin >> color[i];
		color[i]--;
	}
	std::vector<std::vector<int>> edg(n);
	for (int i = 0; i < n - 1; i++) {
		int a, b; std::cin >> a >> b;
		a--; b--;
		edg[a].push_back(b);
		edg[b].push_back(a);
	}

	std::vector<int> pos(2 * n);
	std::vector<std::vector<int>> childs(n);
	auto bfs1 = [&](auto &self, int v, int &id, int prev = -1) -> void {
		pos[2 * v] = id++;
		for (int c : edg[v]) {
			if (c == prev) continue;
			childs[v].push_back(c);
			self(self, c, id, v);
		}
		pos[2 * v + 1] = id++;
	};
	int w = 0;
	bfs1(bfs1, 0, w);
	
	std::vector<int> init(2*n);
	auto bfs2 = [&](auto &self, int v, int prev = -1) -> int {
		int cnt = 0;
		for (int c : edg[v]) {
			if (c == prev) continue;
			cnt += self(self, c, v);
		}
		if (color[v] == 0) {
			init[pos[2 * v]] = init[pos[2 * v + 1]] = -cnt;
			return 0;
		}
		init[pos[2 * v]] = init[pos[2 * v + 1]] = 1;
		return cnt + 1;
	};
	bfs2(bfs2, 0);

	std::vector<std::vector<int>> vs(n);
	for (int v = 0; v < n; v++) vs[color[v]].push_back(v);

	SegmentTree seg(init);
	auto ans = [&](int col) {
		long long ret = (long long)n*(n + 1) / 2;
		for (int v : vs[col]) {
			for (int c : childs[v]) {
				long long k = seg.getval(pos[2 * c], pos[2 * c + 1] + 1) / 2;
				ret -= k*(k + 1) / 2;
			}
		}
		long long k = seg.getval(0, 2 * n) / 2;
		ret -= k*(k + 1) / 2;
		return ret;
	};

	std::cout << ans(0) << std::endl;
	for (int i = 1; i < n; i++) {
		for (int v : vs[i - 1]) {
			seg.update(pos[2 * v], 1);
			seg.update(pos[2 * v + 1], 1);
		}
		for (int v : vs[i]) {
			int cnt = seg.getval(pos[2 * v] + 1, pos[2 * v + 1]) / 2;
			seg.update(pos[2 * v], -cnt);
			seg.update(pos[2 * v + 1], -cnt);
		}
		std::cout << ans(i) << std::endl;
	}
	return 0;
}


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