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

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

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

// std::vector Declaration
template <typename T> vector<T> make_v(size_t a) { return vector<T>(a); }
template <typename T, typename... Ts> auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
}

// std::vector Declaration and Initialization
template <typename T> vector<T> make_vector(size_t a, T x) {
  return vector<T>(a, x);
}
template <typename T, typename U, typename... Ts>
auto make_vector(size_t a, U b, Ts... ts) {
  return vector<decltype(make_vector<T>(b, ts...))>(a,
                                                    make_vector<T>(b, ts...));
}

// std::vector Input
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &e : v)
    is >> e;
  return is;
}

// std::vector Debug
template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  bool a = 1;
  for (auto e : v) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "]";
  return os;
}

// std::array Debug
template <typename T, size_t n>
ostream &operator<<(ostream &os, const array<T, n> &v) {
  os << "[";
  bool a = 1;
  for (auto e : v) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "]";
  return os;
}

// std::deque Debug
template <typename T> ostream &operator<<(ostream &os, const deque<T> &d) {
  os << "[";
  bool a = 1;
  for (auto e : d) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "]";
  return os;
}

// std::pair Debug
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << "(" << p.first << " " << p.second << ")";
  return os;
}

// std::set Debug
template <typename T> ostream &operator<<(ostream &os, const set<T> &st) {
  os << "{";
  bool a = 1;
  for (auto e : st) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "}";
  return os;
}

// std::multiset Debug
template <typename T> ostream &operator<<(ostream &os, const multiset<T> &st) {
  os << "{";
  bool a = 1;
  for (auto e : st) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "}";
  return os;
}

// std::map Debug
template <typename T, typename U>
ostream &operator<<(ostream &os, const map<T, U> &mp) {
  os << "{";
  bool a = 1;
  for (auto e : mp) {
    os << (a ? "" : " ");
    os << e.first << ":" << e.second;
    a = 0;
  }
  os << "}";
  return os;
}

// std::tuple Debug
template <int N, class Tuple> void out(ostream &os, const Tuple &t) {}
template <int N, class Tuple, class H, class... Ts>
void out(ostream &os, const Tuple &t) {
  if (N)
    os << " ";
  os << get<N>(t);
  out<N + 1, Tuple, Ts...>(os, t);
}
template <class... Ts> ostream &operator<<(ostream &os, const tuple<Ts...> &t) {
  os << "(";
  out<0, tuple<Ts...>, Ts...>(os, t);
  os << ")";
  return os;
}

// Debug
#define DUMP(x) cerr << #x << " = " << (x) << endl

// Weighted edge
template <typename T> struct edge {
  int src, to;
  T cost;

  edge() {}
  edge(int to, T cost) : src(-1), to(to), cost(cost) {}
  edge(int src, int to, T cost) : src(src), to(to), cost(cost) {}

  friend ostream &operator<<(ostream &os, const edge &e) {
    return os << "(" << e.src << "->" << e.to << ":" << e.cost << ")";
  }
};

using LL = int64_t;

#define fs first
#define sc second

const int64_t MOD = 1e9 + 7;

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  LL N;
  cin >> N;
  vector<int> c(N);
  for (int i = 0; i < N; ++i) {
    cin >> c[i];
    --c[i];
  }
  vector<vector<int>> G(N);
  for (int i = 0; i < N - 1; ++i) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  // DUMP(G);

  vector<LL> cnt(N), sz(N, 1);
  vector<LL> ans(N);

  function<void(int, int)> dfs = [&](int v, int p) {
    int buff = cnt[c[v]];
    for (int u : G[v]) {
      if (u == p)
        continue;
      cnt[c[v]] = 0;
      dfs(u, v);
      sz[v] += sz[u];
      LL tmp = sz[u] - cnt[c[v]];
      ans[c[v]] += tmp * (tmp + 1) / 2;
    }
    cnt[c[v]] = buff + sz[v];
    // DUMP(v);
    // DUMP(cnt);
    // DUMP(ans);
  };

  dfs(0, -1);
  // DUMP(sz);
  // DUMP(ans);
  // DUMP(cnt);

  for (int i = 0; i < N; ++i) {
    LL tmp = N - cnt[i];
    // DUMP(tmp);
    ans[i] += tmp * (tmp + 1) / 2;
  }

  // DUMP(ans);

  for (int i = 0; i < N; ++i) {
    cout << N * (N + 1) / 2 - ans[i] << endl;
  }

  return 0;
}


Error: TLE

Fixed submission:

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

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

// std::vector Declaration
template <typename T> vector<T> make_v(size_t a) { return vector<T>(a); }
template <typename T, typename... Ts> auto make_v(size_t a, Ts... ts) {
  return vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
}

// std::vector Declaration and Initialization
template <typename T> vector<T> make_vector(size_t a, T x) {
  return vector<T>(a, x);
}
template <typename T, typename U, typename... Ts>
auto make_vector(size_t a, U b, Ts... ts) {
  return vector<decltype(make_vector<T>(b, ts...))>(a,
                                                    make_vector<T>(b, ts...));
}

// std::vector Input
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &e : v)
    is >> e;
  return is;
}

// std::vector Debug
template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  os << "[";
  bool a = 1;
  for (auto e : v) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "]";
  return os;
}

// std::array Debug
template <typename T, size_t n>
ostream &operator<<(ostream &os, const array<T, n> &v) {
  os << "[";
  bool a = 1;
  for (auto e : v) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "]";
  return os;
}

// std::deque Debug
template <typename T> ostream &operator<<(ostream &os, const deque<T> &d) {
  os << "[";
  bool a = 1;
  for (auto e : d) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "]";
  return os;
}

// std::pair Debug
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << "(" << p.first << " " << p.second << ")";
  return os;
}

// std::set Debug
template <typename T> ostream &operator<<(ostream &os, const set<T> &st) {
  os << "{";
  bool a = 1;
  for (auto e : st) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "}";
  return os;
}

// std::multiset Debug
template <typename T> ostream &operator<<(ostream &os, const multiset<T> &st) {
  os << "{";
  bool a = 1;
  for (auto e : st) {
    os << (a ? "" : " ");
    os << e;
    a = 0;
  }
  os << "}";
  return os;
}

// std::map Debug
template <typename T, typename U>
ostream &operator<<(ostream &os, const map<T, U> &mp) {
  os << "{";
  bool a = 1;
  for (auto e : mp) {
    os << (a ? "" : " ");
    os << e.first << ":" << e.second;
    a = 0;
  }
  os << "}";
  return os;
}

// std::tuple Debug
template <int N, class Tuple> void out(ostream &os, const Tuple &t) {}
template <int N, class Tuple, class H, class... Ts>
void out(ostream &os, const Tuple &t) {
  if (N)
    os << " ";
  os << get<N>(t);
  out<N + 1, Tuple, Ts...>(os, t);
}
template <class... Ts> ostream &operator<<(ostream &os, const tuple<Ts...> &t) {
  os << "(";
  out<0, tuple<Ts...>, Ts...>(os, t);
  os << ")";
  return os;
}

// Debug
#define DUMP(x) cerr << #x << " = " << (x) << endl

// Weighted edge
template <typename T> struct edge {
  int src, to;
  T cost;

  edge() {}
  edge(int to, T cost) : src(-1), to(to), cost(cost) {}
  edge(int src, int to, T cost) : src(src), to(to), cost(cost) {}

  friend ostream &operator<<(ostream &os, const edge &e) {
    return os << "(" << e.src << "->" << e.to << ":" << e.cost << ")";
  }
};

using LL = int64_t;

#define fs first
#define sc second

const int64_t MOD = 1e9 + 7;

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);

  LL N;
  cin >> N;
  vector<int> c(N);
  for (int i = 0; i < N; ++i) {
    cin >> c[i];
    --c[i];
  }
  vector<vector<int>> G(N);
  for (int i = 0; i < N - 1; ++i) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  // DUMP(G);

  vector<LL> cnt(N), sz(N, 1);
  vector<LL> ans(N);

  function<void(int, int)> dfs = [&](int v, int p) {
    LL count = 0;
    for (int u : G[v]) {
      if (u == p)
        continue;
      cnt[c[v]] = 0;
      dfs(u, v);
      sz[v] += sz[u];
      count += sz[u] - cnt[c[v]];
      LL tmp = sz[u] - cnt[c[v]];
      ans[c[v]] += tmp * (tmp + 1) / 2;
    }
    // DUMP(v);
    // DUMP(count);
    // DUMP(sz[v] - 1 - count);
    // LL tmp = sz[v] - 1 - count;
    // ans[c[v]] += tmp * (tmp + 1) / 2;
    cnt[c[v]] = sz[v];
  };

  dfs(0, -1);
  DUMP(sz);
  DUMP(ans);
  DUMP(cnt);

  for (int i = 0; i < N; ++i) {
    LL tmp = N - cnt[i];
    DUMP(tmp);
    ans[i] += tmp * (tmp + 1) / 2;
  }

  DUMP(ans);

  for (int i = 0; i < N; ++i) {
    cout << N * (N + 1) / 2 - ans[i] << endl;
  }

  return 0;
}