Problem p02710 (C++)

Buggy submission (status: Runtime Error):
// Coded by tzc_wk
/*
数据不清空，爆零两行泪。
多测不读完，爆零两行泪。
边界不特判，爆零两行泪。
贪心不证明，爆零两行泪。
D P 顺序错，爆零两行泪。
大小少等号，爆零两行泪。
变量不统一，爆零两行泪。
越界不判断，爆零两行泪。
调试不注释，爆零两行泪。
溢出不 l l，爆零两行泪。
忘文件操作，爆零两行泪。
*/
#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define fz(i, a, b) for (int i = a; i <= b; i++)
#define fd(i, a, b) for (int i = a; i >= b; i--)
#define foreach(it, v)                                                         \
  for (__typeof(v.begin()) it = v.begin(); it != v.end(); it++)
#define all(a) a.begin(), a.end()
#define giveup(...) return printf(__VA_ARGS__), 0;
#define fill0(a) memset(a, 0, sizeof(a))
#define fill1(a) memset(a, -1, sizeof(a))
#define fillbig(a) memset(a, 0x3f, sizeof(a))
#define fillsmall(a) memset(a, 0xcf, sizeof(a))
#define mask(a) (1ll << (a))
#define maskx(a, x) ((a) << (x))
#define _bit(a, x) (((a) >> (x)) & 1)
#define _sz(a) ((int)(a).size())
#define filei(a) freopen(a, "r", stdin);
#define fileo(a) freopen(a, "w", stdout);
#define fileio(a)                                                              \
  freopen(a ".in", "r", stdin);                                                \
  freopen(a ".out", "w", stdout)
#define eprintf(...) fprintf(stderr, __VA_ARGS__)
#define put(x) putchar(x)
#define eoln put('\n')
#define space put(' ')
#define y1 y_chenxiaoyan_1
#define y0 y_chenxiaoyan_0
#define int long long
typedef pair<int, int> pii;
inline int read() {
  int x = 0, neg = 1;
  char c = getchar();
  while (!isdigit(c)) {
    if (c == '-')
      neg = -1;
    c = getchar();
  }
  while (isdigit(c))
    x = x * 10 + c - '0', c = getchar();
  return x * neg;
}
inline void print(int x) {
  if (x < 0) {
    putchar('-');
    print(abs(x));
    return;
  }
  if (x <= 9)
    putchar(x + '0');
  else {
    print(x / 10);
    putchar(x % 10 + '0');
  }
}
inline int qpow(int x, int e, int _MOD) {
  int ans = 1;
  while (e) {
    if (e & 1)
      ans = ans * x % _MOD;
    x = x * x % _MOD;
    e >>= 1;
  }
  return ans;
}
int n = read(), a[100005];
int sz[100005];
vector<int> g[100005];
map<int, int> dp[100005];
vector<int> res[100005];
inline void dfs(int x, int f) {
  sz[x] = 1;
  foreach (it, g[x]) {
    int y = *it;
    if (y == f)
      continue;
    dfs(y, x);
    sz[x] += sz[y];
    int t = sz[y];
    if (dp[y].count(a[x]))
      t -= dp[y][a[x]];
    res[a[x]].push_back(t);
    //		cout<<"add "<<a[x]<<" "<<t<<endl;
    if (dp[x].size() < dp[y].size())
      swap(dp[x], dp[y]);
    foreach (it2, dp[y])
      dp[x][it2->first] += it2->second;
  }
  dp[x][a[x]] = sz[x];
  //	cout<<x<<"\t";
  //	fz(i,1,n)	cout<<dp[x][i]<<" ";puts("");
}
signed main() {
  fz(i, 1, n) a[i] = read();
  fz(i, 1, n - 1) {
    int u = read(), v = read();
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs(1, 0);
  fz(i, 1, n) res[i].push_back(n - dp[1][i]);
  fz(i, 1, n) {
    int sum = (n * (n + 1)) / 2;
    foreach (it, res[i]) {
      //			cout<<*it<<" ";
      sum -= (*it) * ((*it) + 1) / 2;
    }
    //		puts("");
    cout << sum << endl;
  }
  return 0;
}

Fixed submission:
// Coded by tzc_wk
/*
数据不清空，爆零两行泪。
多测不读完，爆零两行泪。
边界不特判，爆零两行泪。
贪心不证明，爆零两行泪。
D P 顺序错，爆零两行泪。
大小少等号，爆零两行泪。
变量不统一，爆零两行泪。
越界不判断，爆零两行泪。
调试不注释，爆零两行泪。
溢出不 l l，爆零两行泪。
忘文件操作，爆零两行泪。
*/
#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define fz(i, a, b) for (int i = a; i <= b; i++)
#define fd(i, a, b) for (int i = a; i >= b; i--)
#define foreach(it, v)                                                         \
  for (__typeof(v.begin()) it = v.begin(); it != v.end(); it++)
#define all(a) a.begin(), a.end()
#define giveup(...) return printf(__VA_ARGS__), 0;
#define fill0(a) memset(a, 0, sizeof(a))
#define fill1(a) memset(a, -1, sizeof(a))
#define fillbig(a) memset(a, 0x3f, sizeof(a))
#define fillsmall(a) memset(a, 0xcf, sizeof(a))
#define mask(a) (1ll << (a))
#define maskx(a, x) ((a) << (x))
#define _bit(a, x) (((a) >> (x)) & 1)
#define _sz(a) ((int)(a).size())
#define filei(a) freopen(a, "r", stdin);
#define fileo(a) freopen(a, "w", stdout);
#define fileio(a)                                                              \
  freopen(a ".in", "r", stdin);                                                \
  freopen(a ".out", "w", stdout)
#define eprintf(...) fprintf(stderr, __VA_ARGS__)
#define put(x) putchar(x)
#define eoln put('\n')
#define space put(' ')
#define y1 y_chenxiaoyan_1
#define y0 y_chenxiaoyan_0
#define int long long
typedef pair<int, int> pii;
inline int read() {
  int x = 0, neg = 1;
  char c = getchar();
  while (!isdigit(c)) {
    if (c == '-')
      neg = -1;
    c = getchar();
  }
  while (isdigit(c))
    x = x * 10 + c - '0', c = getchar();
  return x * neg;
}
inline void print(int x) {
  if (x < 0) {
    putchar('-');
    print(abs(x));
    return;
  }
  if (x <= 9)
    putchar(x + '0');
  else {
    print(x / 10);
    putchar(x % 10 + '0');
  }
}
inline int qpow(int x, int e, int _MOD) {
  int ans = 1;
  while (e) {
    if (e & 1)
      ans = ans * x % _MOD;
    x = x * x % _MOD;
    e >>= 1;
  }
  return ans;
}
int n = read(), a[200005];
int sz[200005];
vector<int> g[200005];
map<int, int> dp[200005];
vector<int> res[200005];
inline void dfs(int x, int f) {
  sz[x] = 1;
  foreach (it, g[x]) {
    int y = *it;
    if (y == f)
      continue;
    dfs(y, x);
    sz[x] += sz[y];
    int t = sz[y];
    if (dp[y].count(a[x]))
      t -= dp[y][a[x]];
    res[a[x]].push_back(t);
    //		cout<<"add "<<a[x]<<" "<<t<<endl;
    if (dp[x].size() < dp[y].size())
      swap(dp[x], dp[y]);
    foreach (it2, dp[y])
      dp[x][it2->first] += it2->second;
  }
  dp[x][a[x]] = sz[x];
  //	cout<<x<<"\t";
  //	fz(i,1,n)	cout<<dp[x][i]<<" ";puts("");
}
signed main() {
  fz(i, 1, n) a[i] = read();
  fz(i, 1, n - 1) {
    int u = read(), v = read();
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs(1, 0);
  fz(i, 1, n) res[i].push_back(n - dp[1][i]);
  fz(i, 1, n) {
    int sum = (n * (n + 1)) / 2;
    foreach (it, res[i]) {
      //			cout<<*it<<" ";
      sum -= (*it) * ((*it) + 1) / 2;
    }
    //		puts("");
    cout << sum << endl;
  }
  return 0;
}