alexjercan/bugnet huggingface dataset:
Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
#define LL long long
const int N = 2e5 + 10;
int n, c[N], cnt[N], siz[N];
inline LL cal(int n) {
  // C(n,2) + n
  return 1LL * n * (n + 1) / 2;
}
LL ans[N];
vector<int> G[N];
void dfs(int v, int fa) {
  int tmp = cnt[c[v]];
  siz[v] = 1;
  for (int u : G[v]) {
    if (u == fa)
      continue;
    cnt[c[v]] = 0;
    dfs(u, v);
    siz[v] += siz[u];
    ans[c[v]] += cal(siz[u] - cnt[c[v]]);
  }
  cnt[c[v]] = tmp + siz[v];
}

inline void add(int a, int b) {
  G[a].push_back(b);
  G[b].push_back(a);
}

int main() {
  // scanf("%d", &n);
  std::ios::sync_with_stdio(false);
  cin >> n;
  for (int i = 1; i <= n; i++)
    // scanf("%d", &c[i]);
    cin >> c[i];
  for (int i = 1; i < n; i++) {
    int u, v;
    // scanf("%d%d", &u, &v);
    cin >> u >> v;
    G[u].push_back(v);
    G[v].push_back(u);
    // add(u, v); add(v, u);
  }
  dfs(1, 0);
  for (int i = 1; i <= n; i++) {
    cout << cal(n) - ans[i] - cal(n - cnt[i]) << endl;
  }
  return 0;
}


Error: 0

Fixed submission:


#include <iostream>
#include <vector>
using namespace std;
#define LL long long
const int MAXN = 2e5 + 10;

int N;
int c[MAXN];
vector<int> G[MAXN];
int tot = 0;
int head[MAXN], Next[MAXN], to[MAXN];

// sub[i]: 顶点i下子树顶点的个数
int sub[MAXN];
// cnt[i]: 颜色为i的顶点下子树顶点的个数
int cnt[MAXN];
// ans[i]: 颜色为i的顶点下不经过颜色i的顶点数
LL ans[MAXN];

// 一开始起名count结果和c++ count()冲突，调试了半天, md
inline LL cal(int n) {
  // C(n,2) + n
  return 1LL * n * (n + 1) / 2;
}

void dfs(int v, int fa) {
  int tmp = cnt[c[v]];
  sub[v] = 1;
  for (int i = head[v]; i; i = Next[i]) {
    int u = to[i];
    if (u == fa)
      continue;
    cnt[c[v]] = 0; // 置0重新计数
    dfs(u, v);
    sub[v] += sub[u];
    ans[c[v]] += cal(sub[u] - cnt[c[v]]);
  }
  cnt[c[v]] = tmp + sub[v];
}
inline void add(int a, int b) {
  ++tot;
  Next[tot] = head[a];
  to[tot] = b;
  head[a] = tot;
}

/*
void dfs(int v, int fa) {
        int tmp = cnt[c[v]];
        sub[v] = 1;
        for(int u:G[v]) {
                if(u==fa) continue;
                cnt[c[v]] = 0;
                dfs(u, v);
                sub[v] += sub[u];
                ans[c[v]]+=cal(sub[u]-cnt[c[v]]);
        }
        cnt[c[v]] = tmp + sub[v];
}
* */

int main() {
  std::ios::sync_with_stdio(false);
  cin >> N;
  for (int i = 1; i <= N; i++) {
    cin >> c[i];
  }
  int a, b;
  for (int i = 1; i < N; i++) {
    cin >> a >> b;

    ++tot;
    Next[tot] = head[a];
    to[tot] = b;
    head[a] = tot;

    ++tot;
    Next[tot] = head[b];
    to[tot] = a;
    head[b] = tot;
    // G[a].push_back(b);
    // G[b].push_back(a);
  }
  // dfs1(1, 0);
  dfs(1, 0);
  for (int i = 1; i <= N; i++) {
    cout << cal(N) - ans[i] - cal(N - cnt[i]) << endl;
  }
  return 0;
}
Problem p02710 (C++) — buggy submission, verdict: Time Limit Exceeded

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

template <typename T> void out(T x) {
  cout << x << endl;
  exit(0);
}
#define watch(x) cout << (#x) << " is " << (x) << endl

using ll = long long;

const int maxn = 1e6 + 5;

int n;
int c[maxn];
vector<int> g[maxn];
int siz[maxn];

ll ans[maxn];

void dfs(int at, int p, map<ll, ll> &mp) {
  siz[at] = 1;
  int C = c[at];
  for (int to : g[at]) {
    if (to == p)
      continue;
    map<ll, ll> mp_to;
    dfs(to, at, mp_to);
    siz[at] += siz[to];
    ll middle = siz[to] - mp_to[C];
    ans[C] -= (middle * (middle + 1) / 2);
    if (mp.size() < mp_to.size()) {
      swap(mp, mp_to);
    }
    for (auto pr : mp_to) {
      mp[pr.first] += pr.second;
    }
  }

  mp[C] = siz[at];
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);

  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }

  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }

  for (int i = 1; i <= n; i++) {
    ans[i] = 1ll * n * (n + 1) / 2;
  }

  map<ll, ll> mp;
  dfs(1, 0, mp);

  for (int i = 1; i <= n; i++) {
    ll top = n - mp[i];
    ans[i] -= (top * (top + 1) / 2);
  }

  for (int i = 1; i <= n; i++) {
    cout << ans[i] << "\\n";
  }

  return 0;
}


Error: TLE

Fixed submission:

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

template <typename T> void out(T x) {
  cout << x << endl;
  exit(0);
}
#define watch(x) cout << (#x) << " is " << (x) << endl

using ll = long long;

const int maxn = 1e6 + 5;

int n;
int c[maxn];
vector<int> g[maxn];
int siz[maxn];

ll ans[maxn];

void dfs(int at, int p, map<ll, ll> &mp) {
  siz[at] = 1;
  int C = c[at];
  for (int to : g[at]) {
    if (to == p)
      continue;
    map<ll, ll> mp_to;
    dfs(to, at, mp_to);
    siz[at] += siz[to];
    ll middle = siz[to] - mp_to[C];
    ans[C] -= (middle * (middle + 1) / 2);
    for (auto pr : mp_to) {
      mp[pr.first] += pr.second;
    }
  }

  mp[C] = siz[at];
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);

  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }

  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }

  for (int i = 1; i <= n; i++) {
    ans[i] = 1ll * n * (n + 1) / 2;
  }

  map<ll, ll> mp;
  dfs(1, 0, mp);

  for (int i = 1; i <= n; i++) {
    ll top = n - mp[i];
    ans[i] -= (top * (top + 1) / 2);
  }

  for (int i = 1; i <= n; i++) {
    cout << ans[i] << "\\n";
  }

  return 0;
}
Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
#include <iostream>
using namespace std;
#define _LL long long
#define rep(i, n) for (_LL i = 0; i < (_LL)(n); i++)

#define NMAX 200000
vector<int> b[NMAX];
_LL sum[NMAX];
_LL c[NMAX];
_LL cc[NMAX];
_LL totalcount = 0;

void search(int current, int p) {
  _LL index = 1;
  _LL pcc = cc[c[current]] + totalcount;
  totalcount++;

  for (auto itr = b[current].begin(); itr != b[current].end(); itr++) {
    if (*itr != p) {
      cc[c[current]] = -totalcount;
      search(*itr, current);
      _LL subc = cc[c[current]] + totalcount;
      sum[c[current]] += subc * (subc + 1) / 2;
    }
  }

  cc[c[current]] = pcc - totalcount;
}

int main() {
  _LL n;
  cin >> n;
  rep(i, n) {
    cin >> c[i];
    c[i]--;
  }
  rep(i, n) sum[i] = 0;
  rep(i, n) cc[i] = 0;
  rep(i, n - 1) {
    int u, v;
    cin >> u >> v;
    u--;
    v--;
    b[u].push_back(v);
    b[v].push_back(u);
  }
  search(0, -1);
  rep(i, n) {
    _LL subc = cc[i] + totalcount;
    sum[i] += subc * (subc + 1) / 2;
  }

  _LL total = n * (n + 1) / 2;

  rep(i, n) cout << total - sum[i] << endl;
}

Error: 0

Fixed submission:

#include <bits/stdc++.h>
#include <iostream>
using namespace std;
#define _LL long long
#define rep(i, n) for (_LL i = 0; i < (_LL)(n); i++)

#define NMAX 20000
vector<int> b[NMAX];
_LL sum[NMAX];
_LL c[NMAX];
_LL cc[NMAX];
_LL totalcount = 0;

void search(int current, int p) {
  _LL index = 1;
  _LL pcc = cc[c[current]] + totalcount;
  totalcount++;

  for (auto itr = b[current].begin(); itr != b[current].end(); itr++) {
    if (*itr != p) {
      cc[c[current]] = -totalcount;
      search(*itr, current);
      _LL subc = cc[c[current]] + totalcount;
      sum[c[current]] += subc * (subc + 1) / 2;
    }
  }

  cc[c[current]] = pcc - totalcount;
}

int main() {
  int n;
  cin >> n;
  rep(i, n) {
    cin >> c[i];
    c[i]--;
  }
  rep(i, n) sum[i] = 0;
  rep(i, n) cc[i] = 0;
  rep(i, n - 1) {
    int u, v;
    cin >> u >> v;
    u--;
    v--;
    b[u].push_back(v);
    b[v].push_back(u);
  }
  search(0, -1);
  rep(i, n) {
    _LL subc = cc[i] + totalcount;
    sum[i] += subc * (subc + 1) / 2;
  }

  _LL total = n * (n + 1) / 2;

  rep(i, n) cout << total - sum[i] << endl;
}
Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>

#define long long long int
using namespace std;

// @author: pashka

vector<vector<int>> g;
vector<int> first, last;
vector<int> et;

void dfs(int x, int p) {
  first[x] = et.size();
  last[x] = et.size();
  et.push_back(x);
  for (int y : g[x]) {
    if (y == p)
      continue;
    dfs(y, x);
    last[x] = et.size();
    et.push_back(x);
  }
}

int main() {
  ios::sync_with_stdio(false);

  int n;
  cin >> n;
  vector<int> c(n);
  for (int i = 0; i < n; i++) {
    cin >> c[i];
    c[i]--;
  }
  g.resize(n);
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    x--;
    y--;
    g[x].push_back(y);
    g[y].push_back(x);
  }

  first.resize(n);
  last.resize(n);
  dfs(0, -1);

  vector<vector<int>> pos(n);
  for (int i = 0; i < et.size(); i++) {
    pos[c[et[i]]].push_back(i);
  }

  for (int k = 0; k < n; k++) {
    vector<int> st;
    st.push_back(0);
    int xx = -1;
    vector<int> cc;
    for (int x : pos[k]) {
      st.back() += (x - xx - 2);
      if (x != first[et[x]]) {
        cc.push_back(st.back());
        st.pop_back();
      }
      if (x != last[et[x]]) {
        st.push_back(0);
      }
      xx = x;
    }
    st.back() += (et.size() - xx - 2);
    if (c[0] != k) {
      cc.push_back(st.back());
    }
    long res = (long)n * (n - 1) / 2 + n;
    for (int x : cc) {
      long sz = x / 2 + 1;
      //            cout << sz << " ";
      res -= (sz * (sz - 1) / 2 + sz);
    }
    //        cout << "\\n";
    cout << res << "\\n";
  }

  return 0;
}

Error: -6
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc


Fixed submission:

#include <bits/stdc++.h>

#define long long long int
using namespace std;

// @author: pashka

vector<vector<int>> g;
vector<int> first, last;
vector<int> et;

void dfs(int x, int p) {
  first[x] = et.size();
  et.push_back(x);
  for (int y : g[x]) {
    if (y == p)
      continue;
    dfs(y, x);
    last[x] = et.size();
    et.push_back(x);
  }
}

int main() {
  ios::sync_with_stdio(false);

  int n;
  vector<int> c(n);
  for (int i = 0; i < n; i++) {
    cin >> c[i];
    c[i]--;
  }
  g.resize(n);
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    x--;
    y--;
    g[x].push_back(y);
    g[y].push_back(x);
  }

  first.resize(n);
  last.resize(n);
  dfs(0, -1);

  vector<vector<int>> pos(n);
  for (int i = 0; i < et.size(); i++) {
    pos[c[et[i]]].push_back(i);
  }

  for (int k = 0; k < n; k++) {
    vector<int> st;
    st.push_back(0);
    int xx = -1;
    vector<int> c;
    for (int x : pos[k]) {
      st.back() += (x - xx - 1);
      if (x != first[et[x]]) {
        c.push_back(st.back());
        st.pop_back();
      }
      if (x != last[et[x]]) {
        st.push_back(0);
      }
    }
    c.push_back(st.back());
    long res = (long)n * (n - 1) / 2;
    for (int x : c) {
      long sz = (x + 1) / 2;
      res -= sz * (sz - 1) / 2;
    }
    cout << res << "\\n";
  }

  return 0;
}
Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
#define ll long long
const int N = 2e5 + 10;
const int INF = 0x3f3f3f3f;
const double PI = acos(-1);
const ll mod = 1e9 + 7;
int n, c[N], cnt[N], siz[N];
ll cal(int x) { return 1ll * x * (x - 1) / 2 + x; }
ll ans[N];
vector<int> G[N];
void dfs(int u, int _fa) {
  siz[u] = 1;
  int tmp = cnt[c[u]];
  for (int v : G[u]) {
    if (v == _fa)
      continue;
    cnt[c[u]] = 0;
    dfs(v, u);
    siz[u] += siz[v];
    ans[c[u]] += cal(siz[v] - cnt[c[u]]);
  }
  cnt[c[u]] = tmp + siz[u];
}
int main() {
  scanf("%d", &n);
  for (int i = 1; i <= n; i++)
    scanf("%d", &c[i]);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d%d", &u, &v);
    G[u].push_back(v);
    G[v].push_back(u);
  }
  dfs(1, 0);
  for (int i = 1; i <= n; i++) {
    printf("%lld\\n", cal(n) - ans[i] - cal(n - cnt[i]));
  }
  return 0;
}

Error: 0

Fixed submission:

#include <iostream>
using namespace std;
#define LL long long
const int MAXN = 2e5 + 10;

int N;
int c[MAXN];
int tot = 0;
int head[MAXN], Next[MAXN], to[MAXN], edge[MAXN];

// sub[i]: 顶点i下子树顶点的个数
int sub[MAXN];
// cnt[i]: 颜色为i的顶点下子树顶点的个数
int cnt[MAXN];
// ans[i]: 颜色为i的顶点下不经过颜色i的顶点数
LL ans[MAXN];
inline LL count(int n) {
  // C(n,2) + n
  return 1LL * (n * (n + 1)) / 2;
}
void dfs(int v, int fa) {
  int tmp = cnt[c[v]];
  sub[v] = 1;
  for (int i = head[v]; i; i = Next[i]) {
    int u = to[i];
    if (u == fa)
      continue;
    cnt[c[v]] = 0; // 置0重新计数
    dfs(u, v);
    sub[v] += sub[u];
    ans[c[v]] += count(sub[u] - cnt[c[v]]);
  }
  cnt[c[v]] = tmp + sub[v];
}
void add(int a, int b, int c) {
  ++tot;
  Next[tot] = head[a];
  to[tot] = b;
  edge[tot] = c;
  head[a] = tot;
}
int main() {
  cin >> N;
  for (int i = 1; i <= N; i++) {
    cin >> c[i];
  }
  int a, b;
  for (int i = 1; i < N; i++) {
    cin >> a >> b;
    add(a, b, 1);
    add(b, a, 1);
  }
  // dfs1(1, 0);
  dfs(1, 0);
  for (int i = 1; i <= N; i++) {
    cout << count(N) - ans[i] - count(N - cnt[i]) << endl;
  }
  return 0;
}
Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
using UL = unsigned long;
using ULL = unsigned long long;
using LL = long long;
#define rep(i, n) for (UL i = 0; i < (n); i++)

struct Problem {

  UL N;
  vector<UL> C;
  vector<vector<UL>> E;
  vector<UL> Sz, P, A, I, CN;
  vector<ULL> ans;

  void DFS(UL p) {
    I[p] = A[C[p]];
    for (UL e : E[p]) {
      if (P[p] == e)
        continue;
      P[e] = p;
      DFS(e);

      ULL t = Sz[e] - (A[C[p]] - I[p]);
      ans[C[p]] += t * (t + 1) / 2;
      I[p] += Sz[e];
      A[C[p]] = I[p];
    }
    A[C[p]]++;
    if (p != 0)
      Sz[P[p]] += Sz[p];
  }

  void Solve() {
    cin >> N;
    C.resize(N);
    rep(i, N) {
      cin >> C[i];
      C[i]--;
    }
    E.resize(N);
    rep(i, N - 1) {
      UL u, v;
      cin >> u >> v;
      u--;
      v--;
      E[u].push_back(v);
      E[v].push_back(u);
    }
    Sz.assign(N, 1);
    P.assign(N, -1);
    A.assign(N, 0);
    I.assign(N, 0);
    CN.assign(N, 0);
    ans.assign(N, 0);

    DFS(0);

    rep(i, N) {
      ULL t = N - A[i];
      ans[i] += t * (t + 1) / 2;
      ans[i] = (ULL)N * (N + 1) / 2 - ans[i];
      cout << ans[i] << endl;
    }
  }
};
int main() {
  unique_ptr<Problem> p(new Problem());
  cout << fixed << setprecision(10);
  p->Solve();
  return 0;
}


Error: -11

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using UL = unsigned long;
using ULL = unsigned long long;
using LL = long long;
#define rep(i, n) for (int i = 0; i < (n); i++)

struct Problem {
  UL N;
  vector<UL> C;
  vector<vector<UL>> E;

  vector<ULL> ans;
  vector<UL> F;

  void Concat(map<UL, UL> &d, const map<UL, UL> &s) {
    for (auto i : s)
      d[i.first] += i.second;
  }

  pair<UL, map<UL, UL>> DFS(UL p, map<UL, UL> &back) {
    F[p] = 1;
    map<UL, UL> tmp;
    UL Vcount = 0;
    for (UL e : E[p]) {
      if (F[e])
        continue;
      auto buf = DFS(e, tmp);
      ULL backN = buf.first - buf.second[C[p]];
      ULL forwdN = Vcount - tmp[C[p]];
      ans[C[p]] += backN * forwdN;
      ans[C[p]] -= tmp[C[p]] * buf.second[C[p]];
      Concat(tmp, buf.second);
      Vcount += buf.first;
    }
    Vcount++;
    ULL backN = N - Vcount + 1;
    ULL forwdN = Vcount - tmp[C[p]];
    ans[C[p]] += backN * forwdN;
    tmp[C[p]] = Vcount;
    /*
    cout << "p=" << p << ", Vcount=" << Vcount;
    cout << ", ans={"; rep(i, N) cout << ans[i] << ","; cout << "}";
    cout << ", tmp={"; rep(i, N) cout << tmp[i] << ","; cout << "}";
    cout << ", back={"; rep(i, N) cout << back[i] << ","; cout << "}" << endl;
    */
    return {Vcount, move(tmp)};
  }

  void Solve() {
    scanf("%d", &N);
    C.resize(N);
    rep(i, N) {
      scanf("%d", &C[i]);
      C[i]--;
    }
    E.resize(N);
    rep(i, N - 1) {
      UL U, V;
      scanf("%d%d", &U, &V);
      U--;
      V--;
      E[U].push_back(V);
      E[V].push_back(U);
    }
    ans.resize(N);
    F.resize(N);
    map<UL, UL> r;
    DFS(0, r);
    rep(i, N) cout << ans[i] << endl;
  }
};
int main() {
  unique_ptr<Problem> p(new Problem());
  p->Solve();
  return 0;
}
Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
#include <iostream>
using namespace std;
const long long MOD = 1e9 + 7;
const long long MAXN = 2e5 + 5;
long long n;
vector<long long> v1[MAXN];
long long ans[MAXN];
long long res[MAXN];
long long col[MAXN];

long long sz[MAXN];
long long f(long long x) { return 1LL * x * (x + 1) / 2; }
void dfs(long long curr, long long par) {
  long long tmp = res[col[curr]];
  for (long long x : v1[curr]) {
    if (x != par) {
      res[col[curr]] = 0;
      dfs(x, curr);
      sz[curr] += sz[x];
      ans[col[curr]] -= f(sz[x] - res[col[curr]]);
    }
  }
  res[col[curr]] = tmp + sz[curr];
}
int main() {
  cin >> n;
  for (long long i = 1; i <= n; i++) {
    cin >> col[i];
    sz[i] = 1;
  }
  for (long long i = 1; i < n; i++) {
    long long x, y;
    cin >> x >> y;
    v1[x].push_back(y);
    v1[y].push_back(x);
  }
  dfs(1, 1);
  for (long long i = 1; i <= n; i++) {
    // cout<<col[i]<<endl;
    cout << (ans[i] + f(n) - f(n - res[i])) << endl;
  }
}

Error: 0

Fixed submission:

#include <bits/stdc++.h>
#include <iostream>
using namespace std;
const long long MOD = 1e9 + 7;
const long long MAXN = 5e3 + 5;
long long n;
vector<long long> v1[MAXN];
long long ans[MAXN];
long long res[MAXN];
long long col[MAXN];

long long sz[MAXN];
long long f(long long x) { return 1LL * x * (x + 1) / 2; }
void dfs(long long curr, long long par) {
  long long tmp = res[col[curr]];
  for (long long x : v1[curr]) {
    if (x != par) {
      res[col[curr]] = 0;
      dfs(x, curr);
      sz[curr] += sz[x];
      ans[col[curr]] -= f(sz[x] - res[col[curr]]);
    }
  }
  res[col[curr]] = tmp + sz[curr];
}
int main() {
  cin >> n;
  for (long long i = 1; i <= n; i++) {
    cin >> col[i];
    sz[i] = 1;
  }
  for (long long i = 1; i < n; i++) {
    long long x, y;
    cin >> x >> y;
    v1[x].push_back(y);
    v1[y].push_back(x);
  }
  dfs(1, 1);
  for (long long i = 1; i <= n; i++) {
    // cout<<col[i]<<endl;
    cout << (ans[i] + f(n) - f(n - res[i])) << endl;
  }
}
Problem p02710 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using vl = vector<ll>;
using vvl = vector<vl>;
#define rep(i, n) for (int i = 0; i < (n); ++i)
const int MAX_N = 2e5 + 7;

ll N;
vvl G(MAX_N);
vl cs(MAX_N), tree_size(MAX_N), color_cnts(MAX_N, 0);
vl ans(MAX_N);

ll nc2(ll x) { return x * (x - 1) / 2; }

void calc_size(int v, int par = -1) {
  tree_size[v] = 1;
  for (auto to : G[v]) {
    if (to == par)
      continue;
    calc_size(to, v);
    tree_size[v] += tree_size[to];
  }
}
void dfs(int v, int par = -1) {
  ll prev_cnt = color_cnts[cs[v]];
  for (auto to : G[v]) {
    if (to == par)
      continue;
    ll prev = color_cnts[cs[v]];
    dfs(to, v);
    ll diff = color_cnts[cs[v]] - prev;
    ll cnt = tree_size[to] - diff;
    // cerr << v << "->" << to << ": " << cnt << endl;
    ans[cs[v]] -= nc2(cnt);
  }
  color_cnts[cs[v]] = prev_cnt + tree_size[v];
}

int main() {
  cin >> N;
  ans.assign(MAX_N, nc2(N));
  rep(i, N) cin >> cs[i + 1], ans[cs[i + 1]]++;
  rep(i, N - 1) {
    int u, v;
    cin >> u >> v;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  calc_size(1);
  dfs(1);
  rep(c, N) ans[c + 1] -= nc2(N - color_cnts[c + 1]);
  rep(i, N) cout << ans[i + 1] << "\\n";
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using vl = vector<ll>;
using vvl = vector<vl>;
#define rep(i, n) for (int i = 0; i < (n); ++i)
const int MAX_N = 2e5 + 7;

ll N;
vvl G(MAX_N);
vl cs(MAX_N), tree_size(MAX_N), color_cnts(MAX_N, 0);
vl ans(MAX_N);

ll nc2(ll x) { return x * (x - 1) / 2; }

void calc_size(int v, int par = -1) {
  tree_size[v] = 1;
  for (auto to : G[v]) {
    if (to == par)
      continue;
    calc_size(to, v);
    tree_size[v] += tree_size[to];
  }
}
void dfs(int v, int par = -1) {
  ll prev_cnt = color_cnts[cs[v]];
  for (auto to : G[v]) {
    if (to == par)
      continue;
    ll prev = color_cnts[cs[v]];
    dfs(to, v);
    ll diff = color_cnts[cs[v]] - prev;
    ll cnt = tree_size[to] - diff;
    cerr << v << "->" << to << ": " << cnt << endl;
    ans[cs[v]] -= nc2(cnt);
  }
  color_cnts[cs[v]] = prev_cnt + tree_size[v];
}

int main() {
  cin >> N;
  ans.assign(MAX_N, nc2(N));
  rep(i, N) cin >> cs[i + 1], ans[cs[i + 1]]++;
  rep(i, N - 1) {
    int u, v;
    cin >> u >> v;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  calc_size(1);
  dfs(1);
  rep(c, N) ans[c + 1] -= nc2(N - color_cnts[c + 1]);
  rep(i, N) cout << ans[i + 1] << endl;
  return 0;
}
Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int MAX = 510000;
const int MOD = 1000000007;
int sz[200001], zen[200001], c[200001], ans[200001];
bool used[200001];
vector<int> hen[200001];
int dfssz(int a) {

  sz[a] = 1;
  for (auto &&i : hen[a])
    if (!sz[i])
      sz[a] += dfssz(i);

  return sz[a];
}
void dfs(int a) {
  used[a] = true;
  ans[c[a]] += 1 + zen[c[a]] * (zen[c[a]] - 1) / 2;
  zen[c[a]] -= sz[a];
  ans[c[a]] -= (zen[c[a]] - 1) * zen[c[a]] / 2;
  int tmp = zen[c[a]];
  for (auto &&i : hen[a])
    if (!used[i]) {

      ans[c[a]] -= sz[i] * (sz[i] - 1) / 2;
      zen[c[a]] = sz[i];
      dfs(i);
    }
  zen[c[a]] = tmp;
  return;
}

signed main() {
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }
  int a, b;
  for (int i = 0; i < n - 1; i++) {
    cin >> a >> b;
    hen[a].push_back(b);
    hen[b].push_back(a);
  }
  fill(ans, ans + n + 1, 0);
  fill(sz, sz + n + 1, 0);
  fill(zen, zen + n + 1, n);
  dfssz(1);
  fill(used, used + n + 1, false);
  dfs(1);
  for (int i = 1; i <= n; i++)
    cout << ans[i] << endl;
}


Error: 0

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int MAX = 510000;
const int MOD = 1000000007;
int sz[200001], zen[200001], c[200001], ans[200001];
bool used[200001];
vector<int> hen[200000];
int dfssz(int a) {

  sz[a] = 1;
  for (auto &&i : hen[a])
    if (!sz[i])
      sz[a] += dfssz(i);

  return sz[a];
}
void dfs(int a) {
  used[a] = true;
  ans[c[a]] += 1 + zen[c[a]] * (zen[c[a]] - 1) / 2;
  zen[c[a]] -= sz[a];
  ans[c[a]] -= (zen[c[a]] - 1) * zen[c[a]] / 2;
  int tmp = zen[c[a]];
  for (auto &&i : hen[a])
    if (!used[i]) {

      ans[c[a]] -= sz[i] * (sz[i] - 1) / 2;
      zen[c[a]] = sz[i];
      dfs(i);
    }
  zen[c[a]] = tmp;
  return;
}

signed main() {
  int n;
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> c[i];
  }
  int a, b;
  for (int i = 0; i < n - 1; i++) {
    cin >> a >> b;
    hen[a].push_back(b);
    hen[b].push_back(a);
  }
  fill(ans, ans + n + 1, 0);
  fill(sz, sz + n + 1, 0);
  fill(zen, zen + n + 1, n);
  dfssz(1);
  fill(used, used + n + 1, false);
  dfs(1);
  for (int i = 1; i <= n; i++)
    cout << ans[i] << endl;
}