Language: C++


using namespace std;

class segment_tree {
	vector<int> lc;
	vector<int> rc;
	vector<int> sum;
	int Pool = 0;
	
	public:
	segment_tree(int n) {
		lc.resize(30 * n + 5);
		rc.resize(30 * n + 5);
		sum.resize(30 * n + 5);
	}

	void update(int l, int r, int &x, int p, int v) {
		if(!x) {
			x = ++Pool;
		}

		sum[x] += v;

		if(l == r) {
			return;
		}

		int mid = l + r >> 1;

		if(p <= mid) {
			update(l, mid, lc[x], p, v);
		} else {
			update(mid + 1, r, rc[x], p, v);
		}
	}

	int query(int l, int r, int x, int p) {
		if(!x) {
			return 0;
		}

		if(l == r) {
			return sum[x];
		}

		int mid = l + r >> 1;

		if(p <= mid) {
			return query(l, mid, lc[x], p);
		} else {
			return query(mid + 1, r, rc[x], p);
		}
	}

	int merge(int x, int y) {
		if(!x || !y) {
			return x + y;
		}

		int z = ++Pool;

		lc[z] = merge(lc[x], lc[y]);
		rc[z] = merge(rc[x], rc[y]);
		sum[z] = sum[x] + sum[y];

		return z;
	}
};

int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);

	int n;
	cin >> n;

	vector<int> c;
	c.resize(n);

	for(int &i : c) {
		cin >> i;
		--i;
	}

	vector<vector<int> > G;
	G.resize(n);

	for(int i = 0; i < n - 1; ++i) {
		int u, v;
		cin >> u >> v;

		--u;
		--v;

		G[u].emplace_back(v);
		G[v].emplace_back(u);
	}

	segment_tree t(n);

	vector<int> root;
	root.resize(n);

	vector<long long> ans;
	ans.resize(n);

	vector<int> sz;
	sz.resize(n);

	function<void(int, int)> dfs = [&] (int u, int f) {
		sz[u] = 1;

		int sum = 0;

		for(int v : G[u]) {
			if(v != f) {
				dfs(v, u);
				sz[u] += sz[v];

				int d = t.query(0, n - 1, root[v], c[u]);
				sum += d;
				root[u] = t.merge(root[u], root[v]);
				
				ans[c[u]] += 1LL * (sz[v] - d) * (sz[v] - d + 1) / 2; 
			}
		}

		t.update(0, n - 1, root[u], c[u], sz[u] - sum);
	};

	dfs(0, -1);

	vector<bool> vis;
	vis.resize(n);

	for(int i : c) {
		vis[i] = true;
	}

	for(int i = 0; i < n; ++i) {
		if(vis[i]) {
			int d = n - t.query(0, n - 1, root[0], i);
			ans[i] += 1LL * d * (d + 1) / 2;
			ans[i] = 1LL * n * (n + 1) / 2 - ans[i];
		}
	}

	for(int i : ans) {
		cout << i << '\n';
	}

	return 0;
}


Compiler output:
a.cc:5:9: error: 'vector' does not name a type
    5 |         vector<int> lc;
      |         ^~~~~~
a.cc:6:9: error: 'vector' does not name a type
    6 |         vector<int> rc;
      |         ^~~~~~
a.cc:7:9: error: 'vector' does not name a type
    7 |         vector<int> sum;
      |         ^~~~~~
a.cc: In constructor 'segment_tree::segment_tree(int)':
a.cc:12:17: error: 'lc' was not declared in this scope
   12 |                 lc.resize(30 * n + 5);
      |                 ^~
a.cc:13:17: error: 'rc' was not declared in this scope
   13 |                 rc.resize(30 * n + 5);
      |                 ^~
a.cc:14:17: error: 'sum' was not declared in this scope
   14 |                 sum.resize(30 * n + 5);
      |                 ^~~
a.cc: In member function 'void segment_tree::update(int, int, int&, int, int)':
a.cc:22:17: error: 'sum' was not declared in this scope
   22 |                 sum[x] += v;
      |                 ^~~
a.cc:31:40: error: 'lc' was not declared in this scope; did you mean 'l'?
   31 |                         update(l, mid, lc[x], p, v);
      |                                        ^~
      |                                        l
a.cc:33:44: error: 'rc' was not declared in this scope; did you mean 'r'?
   33 |                         update(mid + 1, r, rc[x], p, v);
      |                                            ^~
      |                                            r
a.cc: In member function 'int segment_tree::query(int, int, int, int)':
a.cc:43:32: error: 'sum' was not declared in this scope
   43 |                         return sum[x];
      |                                ^~~
a.cc:49:46: error: 'lc' was not declared in this scope; did you mean 'l'?
   49 |                         return query(l, mid, lc[x], p);
      |                                              ^~
      |                                              l
a.cc:51:50: error: 'rc' was not declared in this scope; did you mean 'r'?
   51 |                         return query(mid + 1, r, rc[x], p);
      |                                                  ^~
      |                                                  r
a.cc: In member function 'int segment_tree::merge(int, int)':
a.cc:62:17: error: 'lc' was not declared in this scope
   62 |                 lc[z] = merge(lc[x], lc[y]);
      |                 ^~
a.cc:63:17: error: 'rc' was not declared in this scope
   63 |                 rc[z] = merge(rc[x], rc[y]);
      |                 ^~
a.cc:64:17: error: 'sum' was not declared in this scope
   64 |                 sum[z] = sum[x] + sum[y];
      |                 ^~~
a.cc: In function 'int main()':
a.cc:71:9: error: 'ios' has not been declared
   71 |         ios::sync_with_stdio(false);
      |         ^~~
a.cc:72:9: error: 'cin' was not declared in this scope
   72 |         cin.tie(nullptr);
      |         ^~~
a.cc:1:1: note: 'std::cin' is defined in header '<iostream>'; this is probably fixable by adding '#include <iostream>'
  +++ |+#include <iostream>
    1 | 
a.cc:77:9: error: 'vector' was not declared in this scope
   77 |         vector<int> c;
      |         ^~~~~~
a.cc:1:1: note: 'std::vector' is defined in header '<vector>'; this is probably fixable by adding '#include <vector>'
  +++ |+#include <vector>
    1 | 
a.cc:77:16: error: expected primary-expression before 'int'
   77 |         vector<int> c;
      |                ^~~
a.cc:78:9: error: 'c' was not declared in this scope
   78 |         c.resize(n);
      |         ^
a.cc:85:23: error: expected primary-expression before 'int'
   85 |         vector<vector<int> > G;
      |                       ^~~
a.cc:86:9: error: 'G' was not declared in this scope
   86 |         G.resize(n);
      |         ^
a.cc:101:16: error: expected primary-expression before 'int'
  101 |         vector<int> root;
      |                ^~~
a.cc:102:9: error: 'root' was not declared in this scope
  102 |         root.resize(n);
      |         ^~~~
a.cc:104:16: error: expected primary-expression before 'long'
  104 |         vector<long long> ans;
      |                ^~~~
a.cc:105:9: error: 'ans' was not declared in this scope
  105 |         ans.resize(n);
      |         ^~~
a.cc:107:16: error: expected primary-expression before 'int'
  107 |         vector<int> sz;
      |                ^~~
a.cc:108:9: error: 'sz' was not declared in this scope
  108 |         sz.resize(n);
      |         ^~
a.cc:110:9: error: 'function' was not declared in this scope
  110 |         function<void(int, int)> dfs = [&] (int u, int f) {
      |         ^~~~~~~~
a.cc:1:1: note: 'std::function' is defined in header '<functional>'; this is probably fixable by adding '#include <functional>'
  +++ |+#include <functional>
    1 | 
a.cc:110:31: error: expression list treated as compound expression in functional cast [-fpermissive]
  110 |         function<void(int, int)> dfs = [&] (int u, int f) {
      |                               ^
a.cc:110:18: error: expected primary-expression before 'void'
  110 |         function<void(int, int)> dfs = [&] (int u, int f) {
      |                  ^~~~
a.cc:131:9: error: 'dfs' was not declared in this scope
  131 |         dfs(0, -1);
      |         ^~~
a.cc:133:16: error: expected primary-expression before 'bool'
  133 |         vector<bool> vis;
      |                ^~~~
a.cc:134:9: error: 'vis' was not declared in this scope
  134 |         vis.resize(n);
      |         ^~~
a.cc:149:17: error: 'cout' was not declared in this scope
  149 |                 cout << i << '\n';
      |                 ^~~~
a.cc:149:17: note: 'std::cout' is defined in header '<iostream>'; this is probably fixable by adding '#include <iostream>'