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

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <climits>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

template <typename T> struct EdgeT {
  EdgeT(int _u = -1, int _v = -1, T _w = 0) : u(_u), v(_v), w(_w) {}

  int u, v;
  T w;
};

using Edge = EdgeT<int>;

class Solution {
public:
  vector<long long> solve(vector<int> &C, vector<pair<int, int>> &E) {
    int n = C.size();
    init(n);

    for (auto &e : E) {
      int u = e.first;
      int v = e.second;
      graph[u].push_back(edges.size());
      graph[v].push_back(edges.size());
      edges.emplace_back(u, v, 0LL);
    }

    rt = 0;
    clock = 0;
    dfs(rt, -1, clock, 0, 0);

    vector<long long> res(n, 0);

    vector<vector<int>> que(n);
    vector<int> cnts(n, 0);
    function<void(int, int)> calc = [&](int u, int p) {
      int clr = C[u];
      auto &Q = que[clr];

      int sum = 0;
      for (auto e : graph[u]) {
        int v = edges[e].u + edges[e].v - u;
        if (v == p) {
          continue;
        }
        calc(v, u);
        int cnt = size[v];

        int l = 0, r = Q.size();
        while (l < r) {
          int m = (l + r) >> 1;
          if (in[Q[m]] < in[v]) {
            l = m + 1;
          } else {
            r = m;
          }
        }
        for (int i = l; i < Q.size(); ++i) {
          if (isAncestor(v, Q[i])) {
            cnt -= size[Q[i]];
          } else {
            break;
          }
        }
        res[clr] += cnt * 1LL * sum;
        sum += cnt;
      }

      int down = size[u];
      int up = n - size[u] + 1;
      while (!Q.empty()) {
        int v = Q.back();
        // cout << clr << " " << v << " " << u << endl;
        if (isAncestor(u, v)) {
          down -= size[v];
          cnts[clr] -= size[v];
          Q.pop_back();
        } else {
          break;
        }
      }
      Q.push_back(u);
      up -= cnts[clr];
      cnts[clr] += size[u];
      // cout << u << " " << up << " " << down << endl;
      res[clr] += down * 1LL * up;
    };

    calc(rt, -1);

    return res;
  }

private:
  int N;
  int rt;
  int clock;
  vector<Edge> edges;
  vector<vector<int>> graph;
  vector<vector<int>> pars;
  vector<int> depth;
  vector<long long> dis;
  vector<int> size;
  vector<int> in;
  vector<int> out;

  void init(int n) {
    N = n;
    rt = -1;
    clock = 0;
    edges.clear();
    graph.clear();
    graph.resize(n);
    depth.clear();
    depth.resize(n, 0);
    dis.clear();
    dis.resize(n, 0);
    size.clear();
    size.resize(n, 0);
    in.clear();
    in.resize(n, -1);
    out.clear();
    out.resize(n, -1);

    {
      int h = 0;
      while (1 << (h + 1) <= n) {
        ++h;
      }
      pars.clear();
      pars.resize(n, vector<int>(h + 1, -1));
    }
  }

  void dfs(int u, int p, int &clock, int dep, long long d) {
    pars[u][0] = p;
    {
      int h = pars[u].size();
      for (int i = 1; i < h; ++i) {
        if (pars[u][i - 1] == -1) {
          break;
        }
        pars[u][i] = pars[pars[u][i - 1]][i - 1];
      }
    }
    depth[u] = dep;
    dis[u] = d;
    in[u] = ++clock;

    size[u] = 1;
    for (auto e : graph[u]) {
      int v = edges[e].u + edges[e].v - u;
      auto w = edges[e].w;
      if (v == p) {
        continue;
      }
      dfs(v, u, clock, dep + 1, d + w);
      size[u] += size[v];
    }

    out[u] = ++clock;
  }

  int lca(int u, int v) {
    if (u == -1 || v == -1) {
      return u == -1 ? v : u;
    }
    if (u == v) {
      return u;
    }
    if (depth[u] > depth[v]) {
      swap(u, v);
    }
    int h = pars[v].size();
    for (int i = h - 1; i >= 0; --i) {
      int p = pars[v][i];
      if (p != -1 && depth[p] >= depth[u]) {
        v = p;
      }
    }
    if (u == v) {
      return u;
    }
    for (int i = h - 1; i >= 0; --i) {
      if (pars[u][i] != pars[v][i]) {
        u = pars[u][i];
        v = pars[v][i];
      }
    }
    return pars[u][0];
  }

  bool isAncestor(int u, int v) { return in[u] <= in[v] && out[v] <= out[u]; }
};
int main(int argc, char **argv) {
  ios::sync_with_stdio(false);
  cin.tie(0);

  int n;
  cin >> n;
  vector<int> C(n);
  for (int i = 0; i < n; ++i) {
    cin >> C[i];
    --C[i];
  }
  vector<pair<int, int>> edges;
  for (int i = 0; i < n - 1; ++i) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    edges.emplace_back(u, v);
  }
  Solution sol;
  auto res = sol.solve(C, edges);
  for (int i = 0; i < n; ++i) {
    cout << res[i] << "\\n";
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <climits>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace std;

// struct Datum {
//     public:
//         Datum() {
//         }

//         void merge(Datum& other) {
//             auto& B = other.cnts;
//             int n = cnts.size();
//             int m = B.size();

//             for (int i = 0, j = 0; i < n || j < m;) {

//              }
//         }

//         vector<pair<int,int>> cnts;
// };
class Solution {
public:
  vector<long long> solve(vector<int> &C, vector<pair<int, int>> &edges) {
    int n = C.size();
    vector<vector<int>> graph(n);
    for (auto &e : edges) {
      int u = e.first, v = e.second;
      graph[u].push_back(v);
      graph[v].push_back(u);
    }

    vector<int> sizes(n, 0);
    vector<long long> res(n, 0);

    vector<unordered_map<int, int>> cnts(n);

    function<void(int, int)> dfs = [&](int u, int p) {
      int w = -1;
      sizes[u] = 1;

      long long cnt = 0;
      for (auto v : graph[u]) {
        if (v == p) {
          continue;
        }

        dfs(v, u);
        sizes[u] += sizes[v];
        cnt -= sizes[v] * 1LL * (sizes[v] - 1) / 2;
        if (w == -1 || cnts[v].size() > cnts[w].size()) {
          w = v;
        }
      }
      cnt += sizes[u];
      cnt += (sizes[u] - 2) * 1LL * (sizes[u] - 1) / 2;
      res[C[u]] += cnt;

      if (w == -1) {
        cnts[u][C[u]] += 1;
        // res[C[u]] += 1;
        return;
      }

      swap(cnts[w], cnts[u]);

      for (auto &p : cnts[u]) {
        if (p.first != C[u]) {
          res[p.first] += p.second * 1LL * (sizes[u] - sizes[w]);
        }
      }
      cnts[u][C[u]] = sizes[u];

      int clr = C[u];
      for (auto v : graph[u]) {
        if (v == p || v == w) {
          continue;
        }
        for (auto &p : cnts[v]) {
          if (p.first != clr) {
            long long cnt = p.second * 1LL * (sizes[u] - sizes[v]);
            cnt -= p.second * 1LL * cnts[u][p.first];
            res[p.first] += cnt;
          }
        }
        for (auto &p : cnts[v]) {
          if (p.first != clr) {
            cnts[u][p.first] += p.second;
          }
        }
      }

      // cout << u << ":" << endl;
      // for (auto& p : cnts[u]) {
      //     cout << p.first << " " << p.second << endl;
      // }
    };

    dfs(0, -1);
    return res;
  }
};

int main(int argc, char **argv) {
  ios::sync_with_stdio(false);
  cin.tie(0);

  int n;
  cin >> n;
  vector<int> C(n);
  for (int i = 0; i < n; ++i) {
    cin >> C[i];
    --C[i];
  }
  vector<pair<int, int>> edges;
  for (int i = 0; i < n - 1; ++i) {
    int u, v;
    cin >> u >> v;
    --u, --v;
    edges.emplace_back(u, v);
  }
  Solution sol;
  auto res = sol.solve(C, edges);
  for (int i = 0; i < n; ++i) {
    cout << res[i] << "\\n";
  }
  return 0;
}