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

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

template <typename T> struct BIT {
  int n;
  vector<T> d;
  BIT(int n = 0) : n(n), d(n + 1) {}
  void add(int i, T x = 1) {
    for (i++; i <= n; i += i & -i) {
      d[i] += x;
    }
  }
  T sum(int i) {
    T x = 0;
    for (i++; i; i -= i & -i) {
      x += d[i];
    }
    return x;
  }
  T sum(int l, int r) { return sum(r - 1) - sum(l - 1); }
};

#define max_N 200010

int k;
vector<int> edges[max_N];
vector<int> color_vertex[max_N];

int in[max_N];
int out[max_N];

long comb2(int x) { return (long)x * (x - 1) / 2 + x; }

void dfs(int v, int p = -1) {
  in[v] = k++;
  for (auto x : edges[v]) {
    if (x == p)
      continue;
    dfs(x, v);
  }
  out[v] = k;
}

int main() {
  int n;
  scanf("%d", &n);
  for (int i = 0; i < n; i++) {
    int c;
    scanf("%d", &c);
    color_vertex[c - 1].push_back(i);
  }
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    edges[x - 1].push_back(y - 1);
    edges[y - 1].push_back(x - 1);
  }

  long total = comb2(n);

  BIT<int> prefix(n);
  for (int i = 0; i < n; i++)
    prefix.add(i, 1);

  dfs(0);
  // 色をすべて見る
  for (int i = 0; i < n; i++) {
    long ans = total;

    sort(color_vertex[i].begin(), color_vertex[i].end(),
         [&](int a, int b) { return in[a] > in[b]; });

    vector<pair<int, int>> history;
    for (auto v : color_vertex[i]) {
      int vsub_size = 1;
      for (auto u : edges[v]) {
        if (in[u] < in[v])
          continue;
        int usub_size = prefix.sum(in[u], out[u]);
        ans -= comb2(usub_size);
        vsub_size += usub_size;
      }
      prefix.add(in[v], -vsub_size);
      history.emplace_back(in[v], vsub_size);
    }
    ans -= comb2(prefix.sum(0, n));

    for (auto p : history)
      prefix.add(p.first, p.second);

    printf("%ld\\n", ans);
  }
}

Error: TLE

Fixed submission:

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

#define max_N 200010

int k;
vector<int> edges[max_N];
vector<int> color_vertex[max_N];

int in[max_N];
int out[max_N];

long comb2(int x) { return (long)x * (x - 1) / 2 + x; }

void dfs(int v, int p = -1) {
  in[v] = k++;
  for (auto x : edges[v]) {
    if (x == p)
      continue;
    dfs(x, v);
  }
  out[v] = k;
}

int main() {
  int n;
  scanf("%d", &n);
  for (int i = 0; i < n; i++) {
    int c;
    scanf("%d", &c);
    color_vertex[c - 1].push_back(i);
  }
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    edges[x - 1].push_back(y - 1);
    edges[y - 1].push_back(x - 1);
  }

  long total = comb2(n);
  vector<int> prefix(n, 1);

  auto sum_prefix = [&](int a, int b) {
    int sum = 0;
    for (int i = a; i < b; i++)
      sum += prefix[i];
    return sum;
  };

  dfs(0);
  for (int i = 0; i < n; i++) {
    long ans = total;

    sort(color_vertex[i].begin(), color_vertex[i].end(),
         [&](int a, int b) { return in[a] > in[b]; });

    vector<pair<int, int>> history;
    for (auto v : color_vertex[i]) {
      int vsub_size = 1;
      for (auto u : edges[v]) {
        if (in[u] < in[v])
          continue;
        int usub_size = sum_prefix(in[u], out[u]);
        ans -= comb2(usub_size);
        vsub_size += usub_size;
      }
      prefix[in[v]] -= vsub_size;
      history.emplace_back(in[v], vsub_size);
    }
    ans -= comb2(sum_prefix(0, n));

    for (auto p : history)
      prefix[p.first] += p.second;

    printf("%ld\\n", ans);
  }
}