Problem p02710 (C++) — buggy submission, verdict: Time Limit Exceeded

/// __Macro's__
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define II pair<int, int>
#define III pair<int, II>
#define VI vector<int>
#define VII vector<II>
#define VIII vector<III>
#define VVI vector<vector<int>>
#define VVII vector<vector<II>>
#define fr first
#define sc second
#define mkpr make_pair
#define PQ priority_queue
#define pb push_back
#define eb emplace_back
#define all(v) v.begin(), v.end()
#define LINE putc('\\n', stdout)
#define SPACE putc(' ', stdout)
#define TAB putc('\\t', stdout)
#define fillarr(arr, val)                                                      \\
  fill((int *)arr, (int *)arr + (sizeof(arr) / sizeof(int)), val)
#define Log2(x) (31 ^ __builtin_clz(x))
#define inf64 ((ll)1e17)
#define inf 1000000007
#define MOD 1000000007
#define getint ReadInt()
const char IO_MODE = 3; // in-->00<--out (MSB-->LSB) | 0: Norm, 1: Fast
inline ll ReadInt() {
  ll x = 0, s = 1;
  char c = getchar();
  while (c < '0' || c > '9') {
    if (c == '-')
      s = -1;
    c = getchar();
  }
  while (c >= '0' && c <= '9') {
    x = x * 10 + c - '0';
    c = getchar();
  }
  return s * x;
}
inline void WriteInt(ll x) {
  char c[20];
  if (!x) {
    putchar('0');
    return;
  }
  if (x < 0)
    putchar('-'), x = -x;
  int i = 0;
  while (x > 0)
    c[++i] = x % 10, x /= 10;
  while (i)
    putchar(c[i--] + 48);
}
template <typename T> inline void out(T x) {
  if (IO_MODE & 1)
    WriteInt(x);
  else if (typeid(x) == typeid(int))
    printf("%i", x);
  else
    printf("%lld", (ll)x);
}
template <typename T, typename... Args> inline void out(T x, Args... args) {
  out(x);
  SPACE;
  out(args...);
}
template <typename T> inline void in(T &x) {
  if (IO_MODE & 2)
    x = ReadInt();
  else if (typeid(x) == typeid(int))
    scanf("%i", &x);
  else if (typeid(x) == typeid(ll))
    scanf("%lld", &x);
}
template <typename T, typename... Args> inline void in(T &x, Args &...args) {
  in(x);
  in(args...);
}
int tc;
///
/// ---------------------------------------------- (Debug)
/// -----------------------------------------------
#define sim template <class c
#define ris return *this
#define dor > debug &operator<<
#define eni(x)                                                                 \\
  sim > typename enable_if<sizeof dud<c>(0) x 1, debug &>::type operator<<(    \\
            c i) {
sim > struct rge {
  c b, e;
};
sim > rge<c> range(c i, c j) { return rge<c>{i, j}; }
sim > auto dud(c *x) -> decltype(cerr << *x, 0);
sim > char dud(...);
struct debug {
#ifdef LOCAL
  eni(!=) cerr << boolalpha << i;
  ris;
} eni(==) ris << range(begin(i), end(i));
}
sim, class b dor(pair<b, c> d) {
  ris << "(" << d.first << ", " << d.second << ")";
}
sim dor(rge<c> d) {
  *this << "[";
  for (auto it = d.b; it != d.e; ++it)
    *this << ", " + 2 * (it == d.b) << *it;
  ris << "]";
}
#else
  sim dor(const c &) { ris; }
#endif
}
;
vector<char *> tokenizer(const char *args) {
  char *token = new char[111];
  strcpy(token, args);
  token = strtok(token, ", ");
  vector<char *> v({token});
  while (token = strtok(NULL, ", "))
    v.push_back(token);
  return reverse(v.begin(), v.end()), v;
}
void debugg(vector<char *> args) { cerr << "\\b\\b "; }
template <typename Head, typename... Tail>
void debugg(vector<char *> args, Head H, Tail... T) {
  debug() << " [" << args.back() << ": " << H << "] ";
  args.pop_back();
  debugg(args, T...);
}
#define harg(...) #__VA_ARGS__
#ifdef LOCAL
#define dbg(...)                                                               \\
  {                                                                            \\
    debugg(tokenizer(harg(__VA_ARGS__, \\b\\t-- > Line)), __VA_ARGS__,           \\
           __LINE__);                                                          \\
    cerr << endl;                                                              \\
  }
#else
#define dbg(...)                                                               \\
  {}
#endif
/// ------------------------------------------------------------------------------------------------------

#ifdef LOCAL
#define nax 1000
#define qax 1000
#else
#define nax 200200
#define qax 200100
#endif

#define int long long

int n, col[nax], ans[nax], sz[nax];

VI dep[nax << 1];
VI adjList[nax];
VI stk[nax];

void dfs(int u, int par) {
  sz[u] = 1;
  dep[stk[col[u]].back()].pb(u);
  stk[col[u]].pb(u);

  for (auto v : adjList[u])
    if (v != par) {
      dfs(v, u);
      sz[u] += sz[v];

      int x = sz[v];
      while (dep[u].size()) {
        x -= sz[dep[u].back()];
        dep[u].pop_back();
      }

      ans[col[u]] += x * (x + 1) / 2;
    }

  stk[col[u]].pop_back();
}

int32_t main() {
  in(n);
  for (int i = 0; i < n; i++)
    in(col[i]);
  for (int i = 1; i < n; i++) {
    int u = getint - 1, v = getint - 1;
    adjList[u].pb(v);
    adjList[v].pb(u);
  }

  for (int i = 1; i <= n; i++)
    stk[i].pb(i + n);
  dfs(0, 0);

  for (int i = 1; i <= n; i++) {
    int x = n;
    for (auto v : dep[i + n])
      x -= sz[v];
    ans[i] += x * (x + 1) / 2;
  }

  int tot = n * (n + 1) / 2;
  for (int i = 1; i <= n; i++)
    out(tot - ans[i]), LINE;
}


Error: TLE

Fixed submission:

/// __Macro's__
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define II pair<int, int>
#define III pair<int, II>
#define VI vector<int>
#define VII vector<II>
#define VIII vector<III>
#define VVI vector<vector<int>>
#define VVII vector<vector<II>>
#define fr first
#define sc second
#define mkpr make_pair
#define PQ priority_queue
#define pb push_back
#define eb emplace_back
#define all(v) v.begin(), v.end()
#define LINE putc('\\n', stdout)
#define SPACE putc(' ', stdout)
#define TAB putc('\\t', stdout)
#define fillarr(arr, val)                                                      \\
  fill((int *)arr, (int *)arr + (sizeof(arr) / sizeof(int)), val)
#define Log2(x) (31 ^ __builtin_clz(x))
#define inf64 ((ll)1e17)
#define inf 1000000007
#define MOD 1000000007
#define getint ReadInt()
const char IO_MODE = 3; // in-->00<--out (MSB-->LSB) | 0: Norm, 1: Fast
inline ll ReadInt() {
  ll x = 0, s = 1;
  char c = getchar();
  while (c < '0' || c > '9') {
    if (c == '-')
      s = -1;
    c = getchar();
  }
  while (c >= '0' && c <= '9') {
    x = x * 10 + c - '0';
    c = getchar();
  }
  return s * x;
}
inline void WriteInt(ll x) {
  char c[20];
  if (!x) {
    putchar('0');
    return;
  }
  if (x < 0)
    putchar('-'), x = -x;
  int i = 0;
  while (x > 0)
    c[++i] = x % 10, x /= 10;
  while (i)
    putchar(c[i--] + 48);
}
template <typename T> inline void out(T x) {
  if (IO_MODE & 1)
    WriteInt(x);
  else if (typeid(x) == typeid(int))
    printf("%i", x);
  else
    printf("%lld", (ll)x);
}
template <typename T, typename... Args> inline void out(T x, Args... args) {
  out(x);
  SPACE;
  out(args...);
}
template <typename T> inline void in(T &x) {
  if (IO_MODE & 2)
    x = ReadInt();
  else if (typeid(x) == typeid(int))
    scanf("%i", &x);
  else if (typeid(x) == typeid(ll))
    scanf("%lld", &x);
}
template <typename T, typename... Args> inline void in(T &x, Args &...args) {
  in(x);
  in(args...);
}
int tc;
///
/// ---------------------------------------------- (Debug)
/// -----------------------------------------------
#define sim template <class c
#define ris return *this
#define dor > debug &operator<<
#define eni(x)                                                                 \\
  sim > typename enable_if<sizeof dud<c>(0) x 1, debug &>::type operator<<(    \\
            c i) {
sim > struct rge {
  c b, e;
};
sim > rge<c> range(c i, c j) { return rge<c>{i, j}; }
sim > auto dud(c *x) -> decltype(cerr << *x, 0);
sim > char dud(...);
struct debug {
#ifdef LOCAL
  eni(!=) cerr << boolalpha << i;
  ris;
} eni(==) ris << range(begin(i), end(i));
}
sim, class b dor(pair<b, c> d) {
  ris << "(" << d.first << ", " << d.second << ")";
}
sim dor(rge<c> d) {
  *this << "[";
  for (auto it = d.b; it != d.e; ++it)
    *this << ", " + 2 * (it == d.b) << *it;
  ris << "]";
}
#else
  sim dor(const c &) { ris; }
#endif
}
;
vector<char *> tokenizer(const char *args) {
  char *token = new char[111];
  strcpy(token, args);
  token = strtok(token, ", ");
  vector<char *> v({token});
  while (token = strtok(NULL, ", "))
    v.push_back(token);
  return reverse(v.begin(), v.end()), v;
}
void debugg(vector<char *> args) { cerr << "\\b\\b "; }
template <typename Head, typename... Tail>
void debugg(vector<char *> args, Head H, Tail... T) {
  debug() << " [" << args.back() << ": " << H << "] ";
  args.pop_back();
  debugg(args, T...);
}
#define harg(...) #__VA_ARGS__
#ifdef LOCAL
#define dbg(...)                                                               \\
  {                                                                            \\
    debugg(tokenizer(harg(__VA_ARGS__, \\b\\t-- > Line)), __VA_ARGS__,           \\
           __LINE__);                                                          \\
    cerr << endl;                                                              \\
  }
#else
#define dbg(...)                                                               \\
  {}
#endif
/// ------------------------------------------------------------------------------------------------------

#ifdef LOCAL
#define nax 1000
#define qax 1000
#else
#define nax 200200
#define qax 200100
#endif

#define int long long

int n, col[nax], ans[nax], sz[nax];
VI adjList[nax];
map<int, int> mp[nax];

void dfs(int u, int par) {
  sz[u] = 1;
  mp[u][col[u]] = 1;
  for (auto v : adjList[u])
    if (v != par) {
      dfs(v, u);
      if (mp[v].count(col[u])) {
        int y = sz[v] - mp[v][col[u]];
        ans[col[u]] += y * (y + 1) / 2;
      } else {
        int y = sz[v];
        ans[col[u]] += y * (y + 1) / 2;
      }

      if (mp[u].size() > mp[v].size())
        swap(mp[u], mp[v]);

      for (auto x : mp[v])
        mp[u][x.fr] += x.sc;

      sz[u] += sz[v];
      mp[u][col[u]] = sz[u];
    }
}

int32_t main() {
  in(n);
  for (int i = 0; i < n; i++)
    in(col[i]);
  for (int i = 1; i < n; i++) {
    int u = getint - 1, v = getint - 1;
    adjList[u].pb(v);
    adjList[v].pb(u);
  }

  dfs(0, 0);
  int tot = n * (n + 1) / 2;
  for (int i = 1; i <= n; i++) {
    int y = n - mp[0][i];
    ans[i] += y * (y + 1) / 2;
    out(tot - ans[i]), LINE;
  }
}