Language: C

#include <bits/stdc++.h>
using namespace std;

//////////////////////////////   Begin Macros

#define all(x) (x).begin(), (x).end()
#define rep(i, N) for (int i = 0; i < (int)(N); i++)
#define reprev(i, N) for (int i = (int)(N)-1; i >= 0; i--)
#define rep1(i, N) for (int i = 1; i <= (int)(N); i++)
#define rep1rev(i, N) for (int i = (int)(N); i >= 0; i--)
#define forbe(i, b, e) for (int i = (b); i < (e); i++)
#define forberev(i, b, e) for (int i = (e)-1; i >= (b); i--)
#define forfl(i, f, l) for (int i = (f); i <= (l); i++)
#define forflrev(i, f, l) for (int i = (l); i >= (f); i--)
using ll = long long int;
using pii = pair<int, int>;
using pil = pair<int, ll>;
using pli = pair<ll, int>;
using pll = pair<ll, ll>;

template <typename T>
bool chmax(T &m, const T q)
{
  if (m < q)
  {
    m = q;
    return true;
  }
  else
    return false;
}
template <typename T>
bool chmin(T &m, const T q)
{
  if (m > q)
  {
    m = q;
    return true;
  }
  else
    return false;
}
template <typename T1, typename T2>
pair<T1, T2> operator+(const pair<T1, T2> &l, const pair<T1, T2> &r) { return {l.first + r.first, l.second + r.second}; }
template <typename T1, typename T2>
pair<T1, T2> operator-(const pair<T1, T2> &l, const pair<T1, T2> &r) { return {l.first - r.first, l.second - r.second}; }
template <typename T>
pair<T, T> operator*(const pair<T, T> &l, const T &r) { return {l.first * r, l.second * r}; }
template <typename T>
pair<T, T> operator/(const pair<T, T> &l, const T &r) { return {l.first / r, l.second / r}; }
template <typename T>
istream &operator>>(istream &is, vector<T> &vec)
{
  for (auto &v : vec)
    is >> v;
  return is;
}
template <typename T>
ostream &operator<<(ostream &os, const vector<T> &vec)
{
  os << "[";
  for (auto v : vec)
    os << v << ",";
  os << "]";
  return os;
}
template <typename T>
ostream &operator<<(ostream &os, const deque<T> &vec)
{
  os << "deq[";
  for (auto v : vec)
    os << v << ",";
  os << "]";
  return os;
}
template <typename T>
ostream &operator<<(ostream &os, const set<T> &vec)
{
  os << "{";
  for (auto v : vec)
    os << v << ",";
  os << "}";
  return os;
}
template <typename T>
ostream &operator<<(ostream &os, const unordered_set<T> &vec)
{
  os << "{";
  for (auto v : vec)
    os << v << ",";
  os << "}";
  return os;
}
template <typename T>
ostream &operator<<(ostream &os, const multiset<T> &vec)
{
  os << "{";
  for (auto v : vec)
    os << v << ",";
  os << "}";
  return os;
}
template <typename T>
ostream &operator<<(ostream &os, const unordered_multiset<T> &vec)
{
  os << "{";
  for (auto v : vec)
    os << v << ",";
  os << "}";
  return os;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &pa)
{
  os << "(" << pa.first << "," << pa.second << ")";
  return os;
}
template <typename TK, typename TV>
ostream &operator<<(ostream &os, const map<TK, TV> &mp)
{
  os << "{";
  for (auto v : mp)
    os << v.first << "=>" << v.second << ",";
  os << "}";
  return os;
}
template <typename TK, typename TV>
ostream &operator<<(ostream &os, const unordered_map<TK, TV> &mp)
{
  os << "{";
  for (auto v : mp)
    os << v.first << "=>" << v.second << ",";
  os << "}";
  return os;
}

template <typename T>
void reset(vector<T> &v, const T reset_to)
{
  for (auto &x : v)
    x = reset_to;
}
inline int popcount(const unsigned int x) { return __builtin_popcount(x); }
#define dbg(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ") " << __FILE__ << endl;

ll nC2(ll n)
{
  return n * (n - 1) / 2;
}

const int intinf = numeric_limits<int>::max();
const ll llinf = numeric_limits<ll>::max();
const pii udlr[4] = {{1, 0}, {0, 1}, {-1, 0}, {0, -1}};

//////////////////////////////   End Macros

#include "Graph.h"

struct TreeDP
{
  struct DP
  {
    int size_sub;
    vector<int> sizes_sub_different_color;

    DP(int size_sub = 1) : size_sub(size_sub){}; //edit here
  };

  Graph &g;
  vector<int> &c;
  vector<DP> dps;
  vector<int> cnt;
  int total;
  vector<int> sizes_sub_different_color_from_root;
  TreeDP(Graph &g, vector<int> &c) : g(g), c(c), dps(g.size()), cnt(g.size(), 0), total(0), sizes_sub_different_color_from_root(g.size(), 0)
  {
    dfs();
    rep(c, g.size())
    {
      sizes_sub_different_color_from_root[c] = total - cnt[c];
    }
  };

  void dfs(int v = 1, int parent = -1)
  {
    int cnt_bfr = cnt[c[v]];
    for (int child : g[v].ids_dst)
    {
      if (child == parent)
        continue;
      int nvertex_different_colors_bfr = total - cnt[c[v]];
      dfs(child, v);
      int nvertex_different_colors_aft = total - cnt[c[v]];
      int sizes_sub_different_color = nvertex_different_colors_aft - nvertex_different_colors_bfr;
      dps[v].size_sub += dps[child].size_sub;
      dps[v].sizes_sub_different_color.push_back(sizes_sub_different_color);

      // cout << "after a child" << endl;
      // dbg(v);
      // dbg(child);
      // dbg(sizes_sub_different_color);
    }

    cnt[c[v]] = cnt_bfr + dps[v].size_sub;
    total++;
    // cout << "exit a vertex" << endl;
    // dbg(v);
    // dbg(c[v]);
    // dbg(dps[v].size_sub);
    // dbg(cnt[c[v]]);
    // dbg(total);
  }
};

void solve()
{
  int N;
  cin >> N;
  vector<int> cs;
  {
    cs.push_back(0);
    vector<int> cs_(N);
    cin >> cs_;
    cs.insert(cs.end(), all(cs_));
  }
  Graph g(N + 1);
  rep(i, N - 1)
  {
    int a, b;
    cin >> a >> b;
    g.add_edge(a, b);
    g.add_edge(b, a);
  }

  ll npath_total = nC2(N + 1);

  TreeDP tree_dp(g, cs);
  vector<ll> ans(g.size(), npath_total);
  rep1(v, N)
  {
    int c = cs[v];
    for (int size_sub_different_color : tree_dp.dps[v].sizes_sub_different_color)
      ans[c] -= nC2(size_sub_different_color + 1);
  }
  rep1(c, N)
  {
    ans[c] -= nC2(tree_dp.sizes_sub_different_color_from_root[c] + 1);
  }

  rep1(c, N)
  {
    cout << ans[c] << endl;
  }
}

int main()
{
  // srand(time(0));
  cout << fixed << setprecision(15);

  solve();
  return 0;
}


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