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

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

#define LL long long
typedef pair<int, int> pii;

template <typename T> inline void read(T &x) {
  static char _c;
  static bool _f;

  x = 0;
  _f = 0;
  _c = getchar();

  while (_c < '0' || '9' < _c) {
    if (_c == '-')
      _f = true;
    _c = getchar();
  }
  while ('0' <= _c && _c <= '9') {
    x = (x << 1) + (x << 3) + (_c & 15);
    _c = getchar();
  }
  if (_f)
    x = -x;
}

template <typename T> inline void Min(T &x, T y) {
  if (y < x)
    x = y;
}
template <typename T> inline void Max(T &x, T y) {
  if (x < y)
    x = y;
}

#define lowbit(x) ((x) & -(x))
#define lson l, mid, id << 1
#define rson mid + 1, r, id << 1 | 1
#define ls id << 1
#define rs id << 1 | 1
#define MID(l, r) ((l) + (((r) - (l)) >> 1))
#define fi first
#define se second
#define mk make_pair
#define pb push_back

const int INF = 0x3f3f3f3f;
const double pi = (double)acos(-1.0);
const double eps = (double)1e-8;
const int e5 = (int)1e5 + 5;
const int MOD = (int)1e9 + 7;
// const int MOD = (int) 998244353;
// int MOD = 998244353;

inline int sig(double x) { return x < -eps ? -1 : eps < x; }
LL fp(LL a, LL n, LL mod = MOD) {
  if (n < 0)
    a = fp(a, mod - 2, mod), n = -n;
  LL res = 1;
  for (; n; n >>= 1, a = a * a % mod)
    if (n & 1)
      res = res * a % mod;
  return res;
}

struct Mint {
  int x;
  Mint() { x = 0; }
  Mint(int _x) : x(_x) {
    if (x < 0 || x >= MOD)
      x = (x % MOD + MOD) % MOD;
  }
  Mint(LL _x) : x(_x) {
    if (x < 0 || x >= MOD)
      x = (x % MOD + MOD) % MOD;
  }
  Mint operator-() const { return Mint(MOD - x); }
  Mint operator+(const Mint &rhs) const {
    return Mint(x + rhs.x >= MOD ? x + rhs.x - MOD : x + rhs.x);
  }
  Mint operator-(const Mint &rhs) const {
    return Mint(x - rhs.x < 0 ? x - rhs.x + MOD : x - rhs.x);
  }
  Mint operator*(const Mint &rhs) const { return Mint((LL)x * rhs.x % MOD); }
  Mint operator/(const Mint &rhs) const {
    return Mint(x * fp(rhs.x, -1) % MOD);
  }
  Mint &operator+=(const Mint &rhs) {
    x += rhs.x;
    if (x >= MOD)
      x -= MOD;
    return *this;
  }
  Mint &operator*=(const Mint &rhs) {
    x = ((LL)x * rhs.x) % MOD;
    return *this;
  }
  bool operator==(const Mint &rhs) const { return x == rhs.x; }
  bool operator!=(const Mint &rhs) const { return x != rhs.x; }
  friend ostream &operator<<(ostream &out, const Mint &rhs) {
    return out << rhs.x;
  }
  friend istream &operator>>(istream &in, Mint &rhs) { return in >> rhs.x; }
};

const int maxn = (int)2e5 + 20;
const int maxm = (int)1e6 + 20;

int n;
int c[maxn];
vector<int> G[maxn];

int vis[maxn];
int sz[maxn];
LL ans[maxn];

int RT, RT_SZ;

void find_root(int u, int fa, int tsz) {
  sz[u] = 1;
  int mx = 0;
  for (auto v : G[u]) {
    if (v == fa || vis[v])
      continue;
    find_root(v, u, tsz);
    sz[u] += sz[v];
    Max(mx, sz[v]);
  }

  Max(mx, tsz - sz[u]);

  if (mx < RT_SZ) {
    RT_SZ = mx;
    RT = u;
  }
}

void dfs0(int u, int fa) {
  sz[u] = 1;

  for (auto v : G[u]) {
    if (v == fa || vis[v])
      continue;
    dfs0(v, u);
    sz[u] += sz[v];
  }
}

void dfs(int u, int fa, unordered_map<int, int> &mp,
         unordered_map<int, int> &cvis) {
  int cc = c[u];
  if (cvis[cc] == 0)
    mp[cc] += sz[u];

  cvis[cc]++;
  for (auto v : G[u]) {
    if (v == fa || vis[v])
      continue;
    dfs(v, u, mp, cvis);
  }
  cvis[cc]--;
}

void PD(int u, int tsz) {
  RT_SZ = tsz;
  RT = u;
  find_root(u, u, tsz);

  u = RT;
  vis[u] = 1;

  ans[c[u]] += (LL)tsz * (tsz - 1) / 2 + tsz;

  unordered_map<int, int> mp2;

  for (auto v : G[u]) {
    if (vis[v])
      continue;
    dfs0(v, v);

    unordered_map<int, int> mp;
    unordered_map<int, int> cvis;
    dfs(v, v, mp, cvis);

    for (auto pr : mp) {
      int cc = pr.fi, cnt = pr.se;
      if (cc != c[u])
        ans[cc] -= (LL)cnt * (cnt - 1) / 2 + (LL)cnt * (sz[v] - cnt);

      mp2[cc] += cnt;
    }

    ans[c[u]] -= (LL)sz[v] * (sz[v] - 1) / 2 + sz[v];
  }

  for (auto pr : mp2) {
    int cc = pr.fi, cnt = pr.se;
    if (cc != c[u])
      ans[cc] += (LL)cnt * (cnt - 1) / 2 + (LL)cnt * (tsz - cnt);
  }

  for (auto v : G[u]) {
    if (vis[v])
      continue;
    PD(v, sz[v]);
  }
}

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

  for (int i = 1; i < n; i++) {
    int x, y;
    read(x), read(y);
    G[x].pb(y);
    G[y].pb(x);
  }

  PD(1, n);

  for (int i = 1; i <= n; i++)
    printf("%lld\\n", ans[i]);
}

int main(int argc, char **argv) {
#ifdef yukihana0416
  freopen("in.txt", "r", stdin);
// freopen("out.txt", "w", stdout);
#endif // yukihana0416
       //    get_prime();
       //    pre();
  int tc = 1;
  //    read(tc);
  for (int ca = 1; ca <= tc; ca++) {
    //        printf("Case #%d: ", ca);
    work();
  }
  return 0;
}


Error: TLE

Fixed submission:

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

#define LL long long
typedef pair<int, int> pii;

template <typename T> inline void read(T &x) {
  static char _c;
  static bool _f;

  x = 0;
  _f = 0;
  _c = getchar();

  while (_c < '0' || '9' < _c) {
    if (_c == '-')
      _f = true;
    _c = getchar();
  }
  while ('0' <= _c && _c <= '9') {
    x = (x << 1) + (x << 3) + (_c & 15);
    _c = getchar();
  }
  if (_f)
    x = -x;
}

template <typename T> inline void Min(T &x, T y) {
  if (y < x)
    x = y;
}
template <typename T> inline void Max(T &x, T y) {
  if (x < y)
    x = y;
}

#define lowbit(x) ((x) & -(x))
#define lson l, mid, id << 1
#define rson mid + 1, r, id << 1 | 1
#define ls id << 1
#define rs id << 1 | 1
#define MID(l, r) ((l) + (((r) - (l)) >> 1))
#define fi first
#define se second
#define mk make_pair
#define pb push_back

const int INF = 0x3f3f3f3f;
const double pi = (double)acos(-1.0);
const double eps = (double)1e-8;
const int e5 = (int)1e5 + 5;
const int MOD = (int)1e9 + 7;
// const int MOD = (int) 998244353;
// int MOD = 998244353;

inline int sig(double x) { return x < -eps ? -1 : eps < x; }
LL fp(LL a, LL n, LL mod = MOD) {
  if (n < 0)
    a = fp(a, mod - 2, mod), n = -n;
  LL res = 1;
  for (; n; n >>= 1, a = a * a % mod)
    if (n & 1)
      res = res * a % mod;
  return res;
}

struct Mint {
  int x;
  Mint() { x = 0; }
  Mint(int _x) : x(_x) {
    if (x < 0 || x >= MOD)
      x = (x % MOD + MOD) % MOD;
  }
  Mint(LL _x) : x(_x) {
    if (x < 0 || x >= MOD)
      x = (x % MOD + MOD) % MOD;
  }
  Mint operator-() const { return Mint(MOD - x); }
  Mint operator+(const Mint &rhs) const {
    return Mint(x + rhs.x >= MOD ? x + rhs.x - MOD : x + rhs.x);
  }
  Mint operator-(const Mint &rhs) const {
    return Mint(x - rhs.x < 0 ? x - rhs.x + MOD : x - rhs.x);
  }
  Mint operator*(const Mint &rhs) const { return Mint((LL)x * rhs.x % MOD); }
  Mint operator/(const Mint &rhs) const {
    return Mint(x * fp(rhs.x, -1) % MOD);
  }
  Mint &operator+=(const Mint &rhs) {
    x += rhs.x;
    if (x >= MOD)
      x -= MOD;
    return *this;
  }
  Mint &operator*=(const Mint &rhs) {
    x = ((LL)x * rhs.x) % MOD;
    return *this;
  }
  bool operator==(const Mint &rhs) const { return x == rhs.x; }
  bool operator!=(const Mint &rhs) const { return x != rhs.x; }
  friend ostream &operator<<(ostream &out, const Mint &rhs) {
    return out << rhs.x;
  }
  friend istream &operator>>(istream &in, Mint &rhs) { return in >> rhs.x; }
};

const int maxn = (int)2e5 + 20;
const int maxm = (int)1e6 + 20;

int n;
int c[maxn];
vector<int> G[maxn];

int vis[maxn];
int sz[maxn];
LL ans[maxn];
int dep[maxn];

int ffa[20][maxn];

int g[maxn];
int gclk[maxn];
int clk;

int RT, RT_SZ;

int find_fa(int x, int y) {
  if (dep[x] > dep[y])
    swap(x, y);
  int c = dep[y] - dep[x];
  for (int i = 0; i < 20; i++)
    if ((c >> i) & 1)
      y = ffa[i][y];
  if (dep[x] == dep[y])
    return y;

  for (int i = 19; ~i; i--)
    if (ffa[i][x] != ffa[i][y]) {
      x = ffa[i][x];
      y = ffa[i][y];
    }
  return ffa[0][x];
}

void find_root(int u, int fa, int tsz) {
  sz[u] = 1;
  int mx = 0;
  for (auto v : G[u]) {
    if (v == fa || vis[v])
      continue;
    find_root(v, u, tsz);
    sz[u] += sz[v];
    Max(mx, sz[v]);
  }

  Max(mx, tsz - sz[u]);

  if (mx < RT_SZ) {
    RT_SZ = mx;
    RT = u;
  }
}

void dfs0(int u, int fa) {
  ffa[0][u] = fa;
  for (int i = 1; i < 20; i++)
    ffa[i][u] = ffa[i - 1][ffa[i - 1][u]];
  dep[u] = dep[fa] + 1;

  sz[u] = 1;

  for (auto v : G[u]) {
    if (v == fa || vis[v])
      continue;
    dfs0(v, u);
    sz[u] += sz[v];
  }
}

void dfs(int u, int fa, map<int, int> &mp) {
  int cc = c[u];
  if (gclk[cc] != clk) {
    gclk[cc] = clk;
    g[cc] = u;
    mp[cc] += sz[u];
  } else {
    if (find_fa(g[cc], u) != g[cc]) {
      mp[cc] += sz[u];
      g[cc] = u;
    }
  }

  for (auto v : G[u]) {
    if (v == fa || vis[v])
      continue;
    dfs(v, u, mp);
  }
}

void PD(int u, int tsz) {
  RT_SZ = tsz;
  RT = u;
  find_root(u, u, tsz);

  u = RT;
  vis[u] = 1;

  ans[c[u]] += (LL)tsz * (tsz - 1) / 2 + tsz;

  map<int, int> mp2;

  for (auto v : G[u]) {
    if (vis[v])
      continue;
    clk++;
    dfs0(v, v);

    map<int, int> mp;
    dfs(v, v, mp);

    for (auto pr : mp)
      mp2[pr.fi] += pr.se;

    for (auto pr : mp) {
      int cc = pr.fi, cnt = pr.se;
      if (cc != c[u])
        ans[cc] -= (LL)cnt * (cnt - 1) / 2 + (LL)cnt * (sz[v] - cnt);
    }

    ans[c[u]] -= (LL)sz[v] * (sz[v] - 1) / 2 + sz[v];
  }

  for (auto pr : mp2) {
    int cc = pr.fi, cnt = pr.se;
    if (cc != c[u])
      ans[cc] += (LL)cnt * (cnt - 1) / 2 + (LL)cnt * (tsz - cnt);
  }

  for (auto v : G[u]) {
    if (vis[v])
      continue;
    PD(v, sz[v]);
  }
}

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

  for (int i = 1; i < n; i++) {
    int x, y;
    read(x), read(y);
    G[x].pb(y);
    G[y].pb(x);
  }

  PD(1, n);

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

int main(int argc, char **argv) {
#ifdef yukihana0416
  freopen("in.txt", "r", stdin);
// freopen("out.txt", "w", stdout);
#endif // yukihana0416
       //    get_prime();
       //    pre();
  int tc = 1;
  //    read(tc);
  for (int ca = 1; ca <= tc; ca++) {
    //        printf("Case #%d: ", ca);
    work();
  }
  return 0;
}