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

#include <bits/stdc++.h>

using namespace std;

struct edge {
  int from, to;
  edge(int from, int to) : from(from), to(to) {}
};
using node = vector<edge>;
using graph = vector<node>;
void addedge(graph &g, int from, int to, bool bidirectional = false) {
  g[from].emplace_back(edge(from, to));
  if (bidirectional)
    g[to].emplace_back(edge(to, from));
}

class Solution {
public:
  graph g;
  vector<int> c;
  vector<long long> ans;
  vector<int> cnt;
  long long f(int k) { return ((long long)k * (k + 1LL)) / 2; }
  int dfs(int u = 0, int p = -1) {
    int ret = 0;
    for (auto e : g[u]) {
      if (e.to == p)
        continue;
      int p = cnt[0] - cnt[c[u]];
      ret += dfs(e.to, u);
      int q = cnt[0] - cnt[c[u]];
      ans[c[u]] += f(q - p);
      cnt[c[u]] += q - p;
    }
    cnt[0]++;
    cnt[c[u]]++;
    return ret + 1;
  }
  void solve() {
    long long n;
    cin >> n;
    c = vector<int>(n, 0);
    ans = vector<long long>(n + 1, 0ll);
    cnt = vector<int>(n + 1, 0);
    for (int i = 0; i < n; i++)
      cin >> c[i];
    for (int i = 0; i < n; i++)
      c[i];
    g = graph(n);
    for (int i = 0; i < n - 1; i++) {
      int u, v;
      cin >> u >> v;
      addedge(g, u - 1, v - 1, true);
    }
    auto tmp = dfs();
    for (int i = 1; i <= n; i++) {
      ans[i] += f(cnt[0] - cnt[i]);
      cout << f(n) - ans[i] << "\\n";
    }
    cout << "\\n";

    return;
  };
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  Solution solution;
  solution.solve();
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

struct edge {
  int from, to;
  edge(int from, int to) : from(from), to(to) {}
};
using node = vector<edge>;
using graph = vector<node>;
void addedge(graph &g, int from, int to, bool bidirectional = false) {
  g[from].emplace_back(edge(from, to));
  if (bidirectional)
    g[to].emplace_back(edge(to, from));
}

class Solution {
public:
  graph g;
  vector<int> c;
  vector<long long> ans;
  unordered_map<int, int> unite(unordered_map<int, int> &x,
                                unordered_map<int, int> &y) {
    if (y.size() > x.size()) {
      for (auto pp : x) {
        y[pp.first] += pp.second;
      }
      return y;
    } else {
      for (auto pp : y) {
        x[pp.first] += pp.second;
      }
      return x;
    }
  }
  long long f(int k) { return ((long long)k * (k + 1LL)) / 2; }

  unordered_map<int, int> dfs(int u = 0, int p = -1) {
    unordered_map<int, int> ret;
    for (auto e : g[u]) {
      if (e.to == p)
        continue;
      auto tmp = dfs(e.to, u);
      ans[c[u]] += f(tmp[0] - tmp[c[u]]);
      ret = unite(ret, tmp);
    }
    ret[0]++;
    ret[c[u]] = ret[0];
    return ret;
  }
  void solve() {
    long long n;
    cin >> n;
    c = vector<int>(n, 0);
    ans = vector<long long>(n + 1, 0ll);
    for (int i = 0; i < n; i++)
      cin >> c[i];
    for (int i = 0; i < n; i++)
      c[i];
    g = graph(n);
    for (int i = 0; i < n - 1; i++) {
      int u, v;
      cin >> u >> v;
      addedge(g, u - 1, v - 1, true);
    }
    auto tmp = dfs();
    for (int i = 1; i <= n; i++) {
      ans[i] += f(tmp[0] - tmp[i]);
      cout << f(n) - ans[i] << "\\n";
    }
    cout << "\\n";

    return;
  };
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  Solution solution;
  solution.solve();
  return 0;
}