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

#define _USE_MATH_DEFINES
#include <bits/stdc++.h>
using namespace std;
// #define int long long
// #define double long double
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define rep1(i, n) for (int i = 1; i < (int)(n); ++i)
#define repeq(i, n) for (int i = 0; i <= (int)(n); ++i)
#define rep1eq(i, n) for (int i = 1; i <= (int)(n); ++i)
#define rrep(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define rrep1(i, n) for (int i = (int)(n)-1; i > 0; --i)
#define rrepeq(i, n) for (int i = (int)(n); i >= 0; --i)
#define rrep1eq(i, n) for (int i = (int)(n); i > 0; --i)
#define REP(i, a, b) for (int i = (int)(a); i <= (int)(b); ++i)
#define RREP(i, a, b) for (int i = (int)(a); i >= (int)(b); --i)
#define ALL(a) (a).begin(), (a).end()
#define RALL(a) (a).rbegin(), (a).rend()
using ll = long long;
using vi = vector<int>;
using vl = vector<ll>;
using vb = vector<bool>;
template <typename T> using Graph = vector<vector<T>>;
template <typename T> using Spacial = vector<vector<vector<T>>>;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
const int MOD = 1e9 + 7;
const int MOD2 = 998244353;
// const double EPS = 1e-9;
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};
string interval[2] = {" ", "\\n"}; // {" ", "\\n"}

template <typename T> struct is_plural : false_type {};
template <typename T1, typename T2>
struct is_plural<pair<T1, T2>> : true_type {};
template <typename T> struct is_plural<vector<T>> : true_type {};
template <typename T> struct is_plural<complex<T>> : true_type {};

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  return is >> p.first >> p.second;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  return os << p.first << " " << p.second;
}
template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  for (auto itr = vec.begin(); itr != vec.end(); ++itr)
    is >> *itr;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  if (vec.empty())
    return os;
  bool pl = is_plural<T>();
  os << vec.front();
  for (auto itr = ++vec.begin(); itr != vec.end(); ++itr)
    os << interval[pl] << *itr;
  return os;
}

bool CoutYN(bool a, string y = "Yes", string n = "No") {
  cout << (a ? y : n) << "\\n";
  return a;
}

template <typename T1, typename T2> inline bool chmax(T1 &a, T2 b) {
  return a < b && (a = b, true);
}
template <typename T1, typename T2> inline bool chmin(T1 &a, T2 b) {
  return a > b && (a = b, true);
}

long long modpow(int a, long long n, int mod = MOD) {
  long long ret = 1;
  do {
    if (n & 1)
      ret = ret * a % mod;
    a = 1LL * a * a % mod;
  } while (n >>= 1);
  return ret;
}

template <typename T> T GCD(T a, T b) { return b ? GCD(b, a % b) : a; }
template <typename T> T LCM(T a, T b) { return a / GCD(a, b) * b; }

template <typename T1, typename T2>
bool CompareBySecond(pair<T1, T2> a, pair<T1, T2> b) {
  return a.second != b.second ? a.second < b.second : a.first < b.first;
}
// counterclockwise from 12 o'clock direction
template <typename T> bool CompareBySlope(pair<T, T> a, pair<T, T> b) {
  if (a.first * b.first < 0)
    return a.first < b.first;
  if (a.first == 0)
    return a.second >= 0 || b.first > 0;
  if (b.first == 0)
    return b.second < 0 && a.first < 0;
  return a.second * b.first < a.first * b.second;
}

/* -------- <templates end> -------- */

template <typename G> struct LowestCommonAncestor {
  const int LOG;
  vector<int> depth;
  const G &g;
  vector<vector<int>> table;

  LowestCommonAncestor(const G &g)
      : g(g), depth(g.size()), LOG(32 - __builtin_clz(g.size())) {
    table.assign(LOG, vector<int>(g.size(), -1));
  }

  void dfs(int v, int p, int d) {
    table[0][v] = p;
    depth[v] = d;
    for (auto &nv : g[v]) {
      // --------  if <typename G> is not Graph, revise nv (e.g. nv[0]) --------
      if (nv != p)
        dfs(nv, v, d + 1);
    }
  }

  void build() {
    dfs(0, -1, 0);
    for (int k = 0; k + 1 < LOG; ++k) {
      for (int i = 0; i < table[k].size(); ++i) {
        if (table[k][i] == -1)
          table[k + 1][i] = -1;
        else
          table[k + 1][i] = table[k][table[k][i]];
      }
    }
  }

  int query(int u, int v) {
    if (depth[u] > depth[v])
      swap(u, v);
    for (int i = LOG - 1; i >= 0; --i) {
      if (((depth[v] - depth[u]) >> i) & 1)
        v = table[i][v];
    }
    if (u == v)
      return u;
    for (int i = LOG - 1; i >= 0; --i) {
      if (table[i][u] != table[i][v]) {
        u = table[i][u];
        v = table[i][v];
      }
    }
    return table[0][u];
  }
};

Graph<pii> seen;
vi notseen;
vl ans;

int dfs(Graph<int> &tree, LowestCommonAncestor<Graph<int>> &lca, vi &color,
        int v, int p) {
  int ret = 1;
  int c = color[v];
  ans[c]++;
  ll start = 1;

  for (int &nv : tree[v]) {
    if (nv != p) {
      int cnt = dfs(tree, lca, color, nv, v);
      ret += cnt;

      auto itr = seen[c].rbegin();
      while (itr != seen[c].rend() && lca.query(itr->first, nv) == nv) {
        cnt -= itr->second;
        ++itr;
      }

      int pb = itr - seen[c].rbegin();
      rep(i, pb) seen[c].pop_back();

      ans[c] += start * cnt;
      start += cnt;
    }
  }

  notseen[c] -= start;
  ans[c] += start * notseen[c];
  seen[c].emplace_back(v, ret);

  return ret;
}

void solve() {
  int n;
  cin >> n;
  vi color(n);
  cin >> color;
  rep(i, n) color[i]--;
  Graph<int> tree(n);
  rep1(i, n) {
    int a, b;
    cin >> a >> b;
    --a, --b;
    tree[a].emplace_back(b);
    tree[b].emplace_back(a);
  }

  LowestCommonAncestor<Graph<int>> lca(tree);
  lca.build();

  seen.resize(n);
  notseen.assign(n, n);
  ans.assign(n, 0);

  dfs(tree, lca, color, 0, -1);

  interval[0] = "\\n";
  cout << ans << endl;
}

/* -------- <programs end> -------- */

signed main() {
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);
  solve();
  return 0;
}


Error: TLE

Fixed submission:

#define _USE_MATH_DEFINES
#include <bits/stdc++.h>
using namespace std;
// #define int long long
// #define double long double
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define rep1(i, n) for (int i = 1; i < (int)(n); ++i)
#define repeq(i, n) for (int i = 0; i <= (int)(n); ++i)
#define rep1eq(i, n) for (int i = 1; i <= (int)(n); ++i)
#define rrep(i, n) for (int i = (int)(n)-1; i >= 0; --i)
#define rrep1(i, n) for (int i = (int)(n)-1; i > 0; --i)
#define rrepeq(i, n) for (int i = (int)(n); i >= 0; --i)
#define rrep1eq(i, n) for (int i = (int)(n); i > 0; --i)
#define REP(i, a, b) for (int i = (int)(a); i <= (int)(b); ++i)
#define RREP(i, a, b) for (int i = (int)(a); i >= (int)(b); --i)
#define ALL(a) (a).begin(), (a).end()
#define RALL(a) (a).rbegin(), (a).rend()
using ll = long long;
using vi = vector<int>;
using vl = vector<ll>;
using vb = vector<bool>;
template <typename T> using Graph = vector<vector<T>>;
template <typename T> using Spacial = vector<vector<vector<T>>>;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
const int MOD = 1e9 + 7;
const int MOD2 = 998244353;
// const double EPS = 1e-9;
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};
string interval[2] = {" ", "\\n"}; // {" ", "\\n"}

template <typename T> struct is_plural : false_type {};
template <typename T1, typename T2>
struct is_plural<pair<T1, T2>> : true_type {};
template <typename T> struct is_plural<vector<T>> : true_type {};
template <typename T> struct is_plural<complex<T>> : true_type {};

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  return is >> p.first >> p.second;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  return os << p.first << " " << p.second;
}
template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  for (auto itr = vec.begin(); itr != vec.end(); ++itr)
    is >> *itr;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  if (vec.empty())
    return os;
  bool pl = is_plural<T>();
  os << vec.front();
  for (auto itr = ++vec.begin(); itr != vec.end(); ++itr)
    os << interval[pl] << *itr;
  return os;
}

bool CoutYN(bool a, string y = "Yes", string n = "No") {
  cout << (a ? y : n) << "\\n";
  return a;
}

template <typename T1, typename T2> inline bool chmax(T1 &a, T2 b) {
  return a < b && (a = b, true);
}
template <typename T1, typename T2> inline bool chmin(T1 &a, T2 b) {
  return a > b && (a = b, true);
}

long long modpow(int a, long long n, int mod = MOD) {
  long long ret = 1;
  do {
    if (n & 1)
      ret = ret * a % mod;
    a = 1LL * a * a % mod;
  } while (n >>= 1);
  return ret;
}

template <typename T> T GCD(T a, T b) { return b ? GCD(b, a % b) : a; }
template <typename T> T LCM(T a, T b) { return a / GCD(a, b) * b; }

template <typename T1, typename T2>
bool CompareBySecond(pair<T1, T2> a, pair<T1, T2> b) {
  return a.second != b.second ? a.second < b.second : a.first < b.first;
}
// counterclockwise from 12 o'clock direction
template <typename T> bool CompareBySlope(pair<T, T> a, pair<T, T> b) {
  if (a.first * b.first < 0)
    return a.first < b.first;
  if (a.first == 0)
    return a.second >= 0 || b.first > 0;
  if (b.first == 0)
    return b.second < 0 && a.first < 0;
  return a.second * b.first < a.first * b.second;
}

/* -------- <templates end> -------- */

template <typename G> struct LowestCommonAncestor {
  const int LOG;
  vector<int> depth;
  const G &g;
  vector<vector<int>> table;

  LowestCommonAncestor(const G &g)
      : g(g), depth(g.size()), LOG(32 - __builtin_clz(g.size())) {
    table.assign(LOG, vector<int>(g.size(), -1));
  }

  void dfs(int v, int p, int d) {
    table[0][v] = p;
    depth[v] = d;
    for (auto &nv : g[v]) {
      // --------  if <typename G> is not Graph, revise nv (e.g. nv[0]) --------
      if (nv != p)
        dfs(nv, v, d + 1);
    }
  }

  void build() {
    dfs(0, -1, 0);
    for (int k = 0; k + 1 < LOG; ++k) {
      for (int i = 0; i < table[k].size(); ++i) {
        if (table[k][i] == -1)
          table[k + 1][i] = -1;
        else
          table[k + 1][i] = table[k][table[k][i]];
      }
    }
  }

  int query(int u, int v) {
    if (depth[u] > depth[v])
      swap(u, v);
    for (int i = LOG - 1; i >= 0; --i) {
      if (((depth[v] - depth[u]) >> i) & 1)
        v = table[i][v];
    }
    if (u == v)
      return u;
    for (int i = LOG - 1; i >= 0; --i) {
      if (table[i][u] != table[i][v]) {
        u = table[i][u];
        v = table[i][v];
      }
    }
    return table[0][u];
  }
};

Graph<pii> seen;
vi notseen;
vl ans;

int dfs(Graph<int> &tree, LowestCommonAncestor<Graph<int>> &lca, vi &color,
        int v, int p) {
  int ret = 1;
  int c = color[v];
  ans[c]++;
  ll start = 1;

  for (int &nv : tree[v]) {
    if (nv != p) {
      int cnt = dfs(tree, lca, color, nv, v);
      ret += cnt;

      auto itr = seen[c].rbegin();
      while (itr != seen[c].rend() && lca.query(itr->first, nv) == nv) {
        cnt -= itr->second;
        ++itr;
      }

      ans[c] += start * cnt;
      start += cnt;
    }
  }

  notseen[c] -= start;
  ans[c] += start * notseen[c];
  seen[c].emplace_back(v, start);

  return ret;
}

void solve() {
  int n;
  cin >> n;
  vi color(n);
  cin >> color;
  rep(i, n) color[i]--;
  Graph<int> tree(n);
  rep1(i, n) {
    int a, b;
    cin >> a >> b;
    --a, --b;
    tree[a].emplace_back(b);
    tree[b].emplace_back(a);
  }

  LowestCommonAncestor<Graph<int>> lca(tree);
  lca.build();

  seen.resize(n);
  notseen.assign(n, n);
  ans.assign(n, 0);

  dfs(tree, lca, color, 0, -1);

  interval[0] = "\\n";
  cout << ans << endl;
}

/* -------- <programs end> -------- */

signed main() {
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
  cout << fixed << setprecision(10);
  solve();
  return 0;
}