Problem p02710 (C++) — buggy submission, verdict: Runtime Error

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

#define ll long long
#define ld long double
#define pb push_back

#define mo 1000000007
#define inf 1e18
#define rep(i, s, n) for (ll i = s; i < n; i = i + 1)
#define rrep(i, s, n) for (ll i = s; i >= n; i--)
#define adiii                                                                  \\
  \\                               
  ios_base::sync_with_stdio(false);                                            \\
  \\                                      
  cin.tie(NULL);                                                               \\
  cout.tie(NULL);

ll mod(ll n) { return (n % (ll)mo + (ll)mo) % (ll)mo; }
ll gcd(ll a, ll b) {
  if (b == 0)
    return a;
  return gcd(b, a % b);
}

vector<ll> adj[200005];
vector<ll> s(200005, 1);
vector<ll> f(200005, 0);
vector<ll> c(200005);
vector<ll> ans(200005, 0);
vector<bool> vis(200005, false);

ll tc;
void dfs(ll u) {
  vis[u] = true;
  ll eval = f[c[u]];

  for (auto v : adj[u]) {
    if (vis[v] == false) {
      ll bef = tc - f[c[u]];
      dfs(v);
      ll aft = tc - f[c[u]];
      ll an = aft - bef;

      ans[c[u]] -= (an * (an - 1)) / 2;
      s[u] += s[v];
    }
  }
  f[c[u]] = eval + s[u];
  tc++;
}
int main() {
  adiii ll n;
  cin >> n;

  rep(i, 1, n + 1) cin >> c[i];
  rep(i, 0, n - 1) {
    ll u, v;
    cin >> u >> v;
    adj[u].pb(v);
    adj[v].pb(u);
  }
  dfs(1);
  rep(i, 1, n + 1) {
    // cout<<i<<" "<<ans[i]<<endl;
    ll an = n - f[i];
    ans[i] -= (an * (an - 1)) / 2;
    // cout<<i<<"ggg"<<ans[i]<<endl;
  }
  rep(i, 1, n + 1) ans[c[i]]++;
  rep(i, 1, n + 1) { cout << (n * (n - 1)) / 2 + ans[i] << endl; }
}


Error: 0

Fixed submission:

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

#define ll long long
#define ld long double
#define pb push_back

#define mo 1000000007
#define inf 1e18
#define rep(i, s, n) for (ll i = s; i < n; i = i + 1)
#define rrep(i, s, n) for (ll i = s; i >= n; i--)
#define adiii                                                                  \\
  \\                               
  ios_base::sync_with_stdio(false);                                            \\
  \\                                      
  cin.tie(NULL);                                                               \\
  cout.tie(NULL);

ll mod(ll n) { return (n % (ll)mo + (ll)mo) % (ll)mo; }
ll gcd(ll a, ll b) {
  if (b == 0)
    return a;
  return gcd(b, a % b);
}

vector<ll> adj[100005];
vector<ll> s(100005, 1);
vector<ll> f(100005, 0);
vector<ll> c(100005);
vector<ll> ans(100005, 0);
vector<bool> vis(100005, false);
ll eval;
ll tc;
void dfs(ll u) {
  vis[u] = true;
  eval = f[c[u]];

  for (auto v : adj[u]) {
    if (vis[v] == false) {
      ll bef = tc - f[c[u]];
      dfs(v);
      ll aft = tc - f[c[u]];
      ll an = aft - bef;

      ans[c[u]] -= (an * (an - 1)) / 2;
      s[u] += s[v];
    }
  }
  f[c[u]] = eval + s[u];
  tc++;
}
int main() {
  adiii ll n;
  cin >> n;

  rep(i, 1, n + 1) cin >> c[i];
  rep(i, 0, n - 1) {
    ll u, v;
    cin >> u >> v;
    adj[u].pb(v);
    adj[v].pb(u);
  }
  dfs(1);
  rep(i, 1, n + 1) {
    // cout<<i<<" "<<ans[i]<<endl;
    ll an = n - f[i];
    ans[i] -= (an * (an - 1)) / 2;
    // cout<<i<<"ggg"<<ans[i]<<endl;
  }
  rep(i, 1, n + 1) ans[c[i]]++;
  rep(i, 1, n + 1) { cout << (n * (n - 1)) / 2 + ans[i] << endl; }
}