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

// #include <ext/pb_ds/assoc_container.hpp>
// #include <ext/pb_ds/tree_policy.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <unordered_map>
#include <unordered_set>
#include <valarray>
#include <vector>
using namespace std;
#define IOS                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);                                                               \\
  cout.tie(NULL);
#define trace(x) cerr << #x << ": " << x << " " << endl;

typedef int64_t ll;

#define endl '\\n'
#define int ll
ll mod = 1e9 + 7;
ll mod1 = 1e9 + 5;
ll power(ll a, ll b) {
  if (b == 0)
    return 1;
  else if (b % 2 == 0)
    return power(((((a) % mod) * ((a) % mod)) % mod), b / 2) % mod;
  else
    return (((a) % mod) *
            (power(((((a) % mod) * ((a) % mod)) % mod), b / 2) % mod)) %
           mod;
}
inline ll inverse(ll a, ll md) {
  a %= md;
  if (a < 0)
    a += md;
  ll b = md, u = 0, v = 1;
  while (a) {
    ll t = b / a;
    b -= t * a;
    swap(a, b);
    u -= t * v;
    swap(u, v);
  }
  assert(b == 1);
  if (u < 0)
    u += md;
  return u;
}

// using namespace __gnu_pbds;
// #define ordered_set tree<pair<int,int>, null_type,less<pair<int,int>>,
// rb_tree_tag,tree_order_statistics_node_update>
// find_by_order(k)  returns iterator to kth element starting from 0;
// order_of_key(k) returns count of elements strictly smaller than k;
// erase,insert same as normal set
//  mt19937
//  mrand(chrono::high_resolution_clock::now().time_since_epoch().count()); for
//  randomized solution
const int ce = 1e9 + 7;
ll fast_mod(ll input) { return input < ce ? input : input % ce; }
#define pi 3.14159
vector<int> v[200001];
vector<int> eu;
int n;
int a[200001];
int st[400001];
int en[400001];
set<pair<int, int>> s[200001];
set<pair<int, int>> s1[200001];
set<pair<int, int>> temp;
int sub[200001];
int sub1[200001];
int cost[200001];
void dfs(int u, int par) {
  sub[u] = 1;
  eu.push_back(u);
  for (auto i : v[u]) {
    if (i == par)
      continue;
    dfs(i, u);
    sub[u] += sub[i];
  }
  eu.push_back(-u);
}
int ans[200001];
void comp(int u, int par) {
  for (auto i : v[u]) {
    if (i == par)
      continue;
    comp(i, u);
  }
  vector<int> te;
  int x = n - sub[u];
  x -= cost[u];
  te.push_back(x);
  temp.clear();
  for (auto i : v[u]) {
    if (i == par)
      continue;
    auto it = s[a[u]].lower_bound({st[i], -1e9});

    if (it == s[a[u]].end()) {
      te.push_back(sub[i]);
      continue;
    }
    int x = sub[i];
    for (; it != s[a[u]].end(); ++it) {
      int in = (*it).second;
      if (en[i] < st[in]) {
        break;
      }
      temp.insert(*it);
      x -= sub[in];
    }

    te.push_back(x);
  }
  int sum = 0;
  for (auto i : te)
    sum += i;
  sum++;
  for (auto i : te) {
    ans[a[u]] += (sum - i) * i;
    sum -= i;
  }

  ans[a[u]]++;
  for (auto i : temp) {
    s[a[u]].erase({i.first, i.second});
  }
  s[a[u]].insert({st[u], u});
  temp.clear();
}
int32_t main() {
  IOS cin >> n;
  for (int i = 1; i <= n; i++)
    cin >> a[i];
  for (int i = 1; i <= n - 1; i++) {
    int x, y;
    cin >> x >> y;
    v[x].push_back(y);
    v[y].push_back(x);
  }
  memset(ans, 0, sizeof(ans));
  memset(sub1, 0, sizeof(sub1));
  memset(cost, 0, sizeof(cost));
  dfs(1, 0);
  for (int i = 0; i < eu.size(); i++) {
    if (eu[i] > 0) {
      st[eu[i]] = i;
    } else {
      en[abs(eu[i])] = i;
    }
  }
  for (int i = 0; i < eu.size(); i++) {
    if (eu[i] > 0) {
      cost[eu[i]] = sub1[a[eu[i]]];
    }
    if (eu[i] < 0) {
      int u = abs(eu[i]);
      temp.clear();

      auto it1 = s1[a[u]].lower_bound({st[u], -1e9});
      for (; it1 != s1[a[u]].end(); ++it1) {
        int in = (*it1).second;
        if (en[u] < st[in]) {
          break;
        }
        temp.insert(*it1);
        sub1[a[u]] -= sub[in];
      }
      for (auto i : temp)
        s1[a[u]].erase({i.first, i.second});
      s1[a[u]].insert({en[u], u});
      sub1[a[u]] += sub[u];
    }
  }
  comp(1, 0);
  for (int i = 1; i <= n; i++)
    cout << ans[i] << endl;
}


Error: TLE

Fixed submission:

// #include <ext/pb_ds/assoc_container.hpp>
// #include <ext/pb_ds/tree_policy.hpp>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <unordered_map>
#include <unordered_set>
#include <valarray>
#include <vector>
using namespace std;
#define IOS                                                                    \\
  ios_base::sync_with_stdio(false);                                            \\
  cin.tie(NULL);                                                               \\
  cout.tie(NULL);
#define trace(x) cerr << #x << ": " << x << " " << endl;

typedef int64_t ll;

#define endl '\\n'
#define int ll
ll mod = 1e9 + 7;
ll mod1 = 1e9 + 5;
ll power(ll a, ll b) {
  if (b == 0)
    return 1;
  else if (b % 2 == 0)
    return power(((((a) % mod) * ((a) % mod)) % mod), b / 2) % mod;
  else
    return (((a) % mod) *
            (power(((((a) % mod) * ((a) % mod)) % mod), b / 2) % mod)) %
           mod;
}
inline ll inverse(ll a, ll md) {
  a %= md;
  if (a < 0)
    a += md;
  ll b = md, u = 0, v = 1;
  while (a) {
    ll t = b / a;
    b -= t * a;
    swap(a, b);
    u -= t * v;
    swap(u, v);
  }
  assert(b == 1);
  if (u < 0)
    u += md;
  return u;
}

// using namespace __gnu_pbds;
// #define ordered_set tree<pair<int,int>, null_type,less<pair<int,int>>,
// rb_tree_tag,tree_order_statistics_node_update>
// find_by_order(k)  returns iterator to kth element starting from 0;
// order_of_key(k) returns count of elements strictly smaller than k;
// erase,insert same as normal set
//  mt19937
//  mrand(chrono::high_resolution_clock::now().time_since_epoch().count()); for
//  randomized solution
const int ce = 1e9 + 7;
ll fast_mod(ll input) { return input < ce ? input : input % ce; }
#define pi 3.14159
vector<int> v[200001];
vector<int> eu;
int n;
int a[200001];
int st[400001];
int en[400001];
set<pair<int, int>> s[200001];
set<pair<int, int>> s1[200001];
set<pair<int, int>> temp;
int sub[200001];
void dfs(int u, int par) {
  eu.push_back(u);
  for (auto i : v[u]) {
    if (i == par)
      continue;
    dfs(i, u);
  }
  eu.push_back(-u);
}
int ans[200001];
void comp(int u, int par) {
  sub[u] = 1;
  for (auto i : v[u]) {
    if (i == par)
      continue;
    comp(i, u);
    sub[u] += sub[i];
  }
  vector<int> te;
  int x = n - sub[u];
  auto it1 = s1[a[u]].lower_bound({st[u], -1e9});
  if (it1 != s1[a[u]].begin()) {
    it1--;
    while (1) {
      x -= sub[(*it1).second];
      if (it1 == s1[a[u]].begin())
        break;
      it1--;
    }
  }
  te.push_back(x);
  temp.clear();
  for (auto i : v[u]) {
    if (i == par)
      continue;
    auto it = s[a[u]].lower_bound({st[i], -1e9});

    if (it == s[a[u]].end()) {
      te.push_back(sub[i]);
      continue;
    }
    int x = sub[i];
    for (; it != s[a[u]].end(); ++it) {
      int in = (*it).second;
      if (en[i] < st[in]) {
        break;
      }
      temp.insert(*it);
      x -= sub[in];
    }

    te.push_back(x);
  }
  int sum = 0;
  for (auto i : te)
    sum += i;
  sum++;
  for (auto i : te) {
    ans[a[u]] += (sum - i) * i;
    sum -= i;
  }

  ans[a[u]]++;
  for (auto i : temp) {
    s[a[u]].erase({i.first, i.second});
  }
  s[a[u]].insert({st[u], u});
  temp.clear();
  it1 = s1[a[u]].lower_bound({st[u], -1e9});
  for (; it1 != s1[a[u]].end(); ++it1) {
    int in = (*it1).second;
    if (en[u] < st[in]) {
      break;
    }
    temp.insert(*it1);
  }
  for (auto i : temp)
    s1[a[u]].erase({i.first, i.second});
  s1[a[u]].insert({en[u], u});
}
int32_t main() {
  IOS

          cin >>
      n;
  for (int i = 1; i <= n; i++)
    cin >> a[i];
  for (int i = 1; i <= n - 1; i++) {
    int x, y;
    cin >> x >> y;
    v[x].push_back(y);
    v[y].push_back(x);
  }
  memset(ans, 0, sizeof(ans));

  dfs(1, 0);
  for (int i = 0; i < eu.size(); i++) {
    if (eu[i] > 0) {
      st[eu[i]] = i;
    } else {
      en[abs(eu[i])] = i;
    }
  }
  comp(1, 0);
  for (int i = 1; i <= n; i++)
    cout << ans[i] << endl;
}