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

// #include "pch.h"
#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <complex>
#include <ctime>
#include <fstream>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>

typedef long long ll;
typedef std::pair<int, int> Pii;
typedef std::pair<long long, long long> Pll;
typedef std::pair<double, double> Pdd;

#define rip(i, n, s) for (int i = (s); i < (int)(n); i++)
#define mapa make_pair
#define all(a) a.begin(), a.end()
#define MM << " " <<

template <typename T> using MaxHeap = std::priority_queue<T>;
template <typename T>
using MinHeap = std::priority_queue<T, std::vector<T>, std::greater<T>>;

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

template <typename T> void vdeb(std::vector<T> &da) {
  for (int i = 0; i < da.size(); i++) {
    if (i == da.size() - 1)
      std::cout << da[i];
    else
      std::cout << da[i] << ' ';
  }
  std::cout << std::endl;
}
template <typename T> void vdeb(std::vector<std::vector<T>> &da) {
  for (int i = 0; i < da.size(); i++)
    vdeb(da[i]);
  std::cout << std::endl;
}

using namespace std;

int dfs1(vector<vector<int>> &da, vector<int> &c, vector<bool> &used,
         vector<ll> &co, int now) {
  rip(i, da[now].size(), 0) {
    if (used[da[now][i]])
      continue;
    used[da[now][i]] = true;
    co[now] += dfs1(da, c, used, co, da[now][i]);
  }
  co[now]++;
  return co[now];
}

ll cal(ll x) { return x * (x + 1) / 2; }

void dfs2(vector<vector<int>> &da, vector<int> &c, vector<bool> &used,
          vector<ll> &ans, vector<ll> &lam, vector<ll> &co, int now) {
  ll tmp = lam[c[now]];
  lam[c[now]] = 0;
  rip(i, da[now].size(), 0) {
    if (used[da[now][i]])
      continue;
    used[da[now][i]] = true;
    dfs2(da, c, used, ans, lam, co, da[now][i]);
    ans[c[now]] += cal(co[da[now][i]] - lam[c[now]]);
    // cout  << co[da[now][i]] MM lam[c[now]] << endl;
    lam[c[now]] = 0;
  }
  lam[c[now]] = tmp + co[now];
  return;
}

int main() {
  ll n;
  cin >> n;
  vector<int> c(n);
  rip(i, n, 0) cin >> c[i], --c[i];
  vector<vector<int>> da(n, vector<int>(0));
  rip(i, n - 1, 0) {
    int a, b;
    cin >> a >> b;
    --a;
    --b;
    da[a].push_back(b);
    da[b].push_back(a);
  }
  vector<bool> used(n, false);
  vector<ll> co(n);
  vector<ll> ans(n), lam(n);
  used[0] = true;
  dfs1(da, c, used, co, 0);
  rip(i, n, 1) used[i] = false;
  // used[0] =true;
  dfs2(da, c, used, ans, lam, co, 0);
  // vdeb(lam);
  // vdeb(co);
  // vdeb(ans);
  rip(i, n, 0) { ans[i] += cal(n - lam[i]); }
  rip(i, n, 0) { printf("%lld\\n", n * (n + 1) / 2 - ans[i]); }
}

Error: TLE

Fixed submission:

// #include "pch.h"
#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <complex>
#include <ctime>
#include <fstream>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>

typedef long long ll;
typedef std::pair<int, int> Pii;
typedef std::pair<long long, long long> Pll;
typedef std::pair<double, double> Pdd;

#define rip(i, n, s) for (int i = (s); i < (int)(n); i++)
#define mapa make_pair
#define all(a) a.begin(), a.end()
#define MM << " " <<

template <typename T> using MaxHeap = std::priority_queue<T>;
template <typename T>
using MinHeap = std::priority_queue<T, std::vector<T>, std::greater<T>>;

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

template <typename T> void vdeb(std::vector<T> &da) {
  for (int i = 0; i < da.size(); i++) {
    if (i == da.size() - 1)
      std::cout << da[i];
    else
      std::cout << da[i] << ' ';
  }
  std::cout << std::endl;
}
template <typename T> void vdeb(std::vector<std::vector<T>> &da) {
  for (int i = 0; i < da.size(); i++)
    vdeb(da[i]);
  std::cout << std::endl;
}

using namespace std;

int dfs1(vector<vector<int>> &da, vector<int> &c, vector<bool> &used,
         vector<ll> &co, int now) {
  rip(i, da[now].size(), 0) {
    if (used[da[now][i]])
      continue;
    used[da[now][i]] = true;
    co[now] += dfs1(da, c, used, co, da[now][i]);
  }
  co[now]++;
  return co[now];
}

ll cal(ll x) { return x * (x + 1) / 2; }

void dfs2(vector<vector<int>> &da, vector<int> c, vector<bool> &used,
          vector<ll> &ans, vector<ll> &lam, vector<ll> &co, int now) {
  ll tmp = lam[c[now]];
  lam[c[now]] = 0;
  rip(i, da[now].size(), 0) {
    if (used[da[now][i]])
      continue;
    used[da[now][i]] = true;
    dfs2(da, c, used, ans, lam, co, da[now][i]);
    ans[c[now]] += cal(co[da[now][i]] - lam[c[now]]);
    // cout  << co[da[now][i]] MM lam[c[now]] << endl;
    lam[c[now]] = 0;
  }
  lam[c[now]] = tmp + co[now];
  return;
}

int main() {
  int n;
  cin >> n;
  vector<int> c(n);
  rip(i, n, 0) cin >> c[i], --c[i];
  vector<vector<int>> da(n, vector<int>(0));
  rip(i, n - 1, 0) {
    int a, b;
    cin >> a >> b;
    --a;
    --b;
    da[a].push_back(b);
    da[b].push_back(a);
  }
  vector<bool> used(n, false);
  vector<ll> co(n);
  vector<ll> ans(n), lam(n);
  used[0] = true;
  dfs1(da, c, used, co, 0);
  rip(i, n, 1) used[i] = false;
  // used[0] =true;
  dfs2(da, c, used, ans, lam, co, 0);
  // vdeb(lam);
  // vdeb(co);
  // vdeb(ans);
  rip(i, n, 0) { ans[i] += cal(n - lam[i]); }
  rip(i, n, 0) { printf("%lld\\n", n * (n + 1) / 2 - ans[i]); }
}