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

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>

#pragma GCC optimize("O3")

using namespace std;
#define PB push_back
#define MP make_pair
#define LL long long
// #define int LL
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define RE(i, n) FOR(i, 1, n)
#define REP(i, n) FOR(i, 0, (int)(n)-1)
#define R(i, n) REP(i, n)
#define VI vector<int>
#define PII pair<int, int>
#define LD long double
#define FI first
#define SE second
#define st FI
#define nd SE
#define ALL(x) (x).begin(), (x).end()
#define SZ(x) ((int)(x).size())

#define unordered_map __fast_unordered_map
template <class Key, class Value, class Hash = std::hash<Key>>
using unordered_map = __gnu_pbds::gp_hash_table<Key, Value, Hash>;

template <class C> void mini(C &a4, C b4) { a4 = min(a4, b4); }
template <class C> void maxi(C &a4, C b4) { a4 = max(a4, b4); }

template <class TH> void _dbg(const char *sdbg, TH h) {
  cerr << sdbg << '=' << h << endl;
}
template <class TH, class... TA> void _dbg(const char *sdbg, TH h, TA... a) {
  while (*sdbg != ',')
    cerr << *sdbg++;
  cerr << '=' << h << ',';
  _dbg(sdbg + 1, a...);
}

template <class T> ostream &operator<<(ostream &os, vector<T> V) {
  os << "[";
  for (auto vv : V)
    os << vv << ",";
  return os << "]";
}
template <class L, class R> ostream &operator<<(ostream &os, pair<L, R> P) {
  return os << "(" << P.st << "," << P.nd << ")";
}
template <class L, class R, class X>
ostream &operator<<(ostream &os, tuple<L, R, X> P) {
  return os << "(" << get<0>(P) << "," << get<1>(P) << "," << get<2>(P) << ")";
}

#ifdef LOCAL
#define debug(...) _dbg(#__VA_ARGS__, __VA_ARGS__)
#else
#define debug(...) (__VA_ARGS__)
#define cerr                                                                   \\
  if (0)                                                                       \\
  cout
#endif

const int MaxN = 2e5 + 100;

int groups[MaxN];

struct BestInfo {
  LL dist;
  // int subtree_idx;
  int vert;

  bool operator<(const BestInfo &other) const { return dist < other.dist; }

  bool SameType(const BestInfo &other) const {
    return groups[vert] == groups[other.vert];
  }

  friend ostream &operator<<(ostream &os, const BestInfo &bi) {
    return os << "(" << bi.dist << "," << bi.vert << ")";
  }
};

vector<PII> adj[MaxN];
int weights[MaxN];
int N;
bool is_avail[MaxN];
int tree_size[MaxN];
int layer_idx[MaxN];
vector<int> conn_verts;

void DfsSizes(int v, int p) {
  tree_size[v] = 1;
  conn_verts.push_back(v);
  for (auto &E : adj[v]) {
    const int s = E.st;
    if (s == p) {
      continue;
    }
    if (!is_avail[s]) {
      continue;
    }
    DfsSizes(s, v);
    tree_size[v] += tree_size[s];
  }
}

void Centroids(int v, int layer) {
  conn_verts.clear();
  DfsSizes(v, v);

  int centroid = v;
  for (int s : conn_verts)
    if (2 * tree_size[s] >= tree_size[v] && tree_size[s] < tree_size[centroid])
      centroid = s;

  is_avail[centroid] = false;
  layer_idx[centroid] = layer;
  debug(centroid, layer);

  for (auto &E : adj[centroid])
    if (is_avail[E.st])
      Centroids(E.st, layer + 1);
}

const int kMaxBests = 4;
int subtree_idx;
int centr_cur;
LL dists_from_centroid[MaxN];
int vert_ptr[MaxN];
LL vert_ptr_dist[MaxN];
BestInfo best_non_centroid[MaxN][2];
BestInfo best_centroid[MaxN];

void DfsDists(int v, int p, LL cur_dist) {
  dists_from_centroid[v] = cur_dist;
  conn_verts.push_back(v);

  if (v != p) {
    BestInfo my_info{cur_dist + weights[v], v};

    if (groups[v] != groups[centr_cur]) {
      if (my_info < best_non_centroid[centr_cur][0]) {
        best_non_centroid[centr_cur][1] = best_non_centroid[centr_cur][0];
        best_non_centroid[centr_cur][0] = my_info;
      } else if (my_info < best_non_centroid[centr_cur][1]) {
        best_non_centroid[centr_cur][1] = my_info;
      }
    } else {
      mini(best_centroid[centr_cur], BestInfo{cur_dist + weights[v], v});
    }
  }

  for (auto &E : adj[v]) {
    const int s = E.st;
    if (s == p || layer_idx[s] < layer_idx[centr_cur]) {
      continue;
    }
    if (v == p) {
      subtree_idx = s;
    }
    DfsDists(s, v, cur_dist + E.nd);
  }
}

void CentroidGroups() {
  fill_n(vert_ptr_dist, N + 1, (LL)1e18);

  for (int centroid = 1; centroid <= N; ++centroid) {
    subtree_idx = -1;
    centr_cur = centroid;
    conn_verts.clear();
    best_non_centroid[centroid][0] = BestInfo{(LL)1e18, -1};
    best_non_centroid[centroid][1] = BestInfo{(LL)1e18, -1};
    best_centroid[centroid] = BestInfo{weights[centroid], centroid};
    DfsDists(centroid, centroid, 0);
    debug(conn_verts);

    for (int v : conn_verts) {
      for (int idx : {0, 1}) {
        const auto &best = best_non_centroid[centroid][idx];
        if (best.vert != -1 && groups[v] != groups[best.vert]) {
          const LL new_dist = best.dist + dists_from_centroid[v] + weights[v];
          if (new_dist < vert_ptr_dist[v]) {
            vert_ptr[v] = best.vert;
            vert_ptr_dist[v] = new_dist;
          }
        }
      }

      if (groups[centroid] != groups[v]) {
        const LL centr_dist =
            dists_from_centroid[v] + weights[v] + weights[centroid];
        const LL group_dist =
            dists_from_centroid[v] + best_centroid[centroid].dist + weights[v];
        if (group_dist < vert_ptr_dist[v]) {
          vert_ptr[v] = best_centroid[centroid].vert;
          vert_ptr_dist[v] = group_dist;
        }
        if (centr_dist < vert_ptr_dist[centroid]) {
          vert_ptr[centroid] = v;
          vert_ptr_dist[centroid] = centr_dist;
        }
      }
    }
  }
}

int fu_parent[MaxN];

int FuFind(int v) {
  if (fu_parent[v] == v)
    return v;
  return fu_parent[v] = FuFind(fu_parent[v]);
}

bool FuUnion(int a, int b) {
  a = FuFind(a);
  b = FuFind(b);
  if (a != b) {
    fu_parent[a] = b;
    return true;
  }
  return false;
}

int32_t main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout << fixed << setprecision(11);
  cerr << fixed << setprecision(6);

  cin >> N;
  for (int i = 1; i <= N; ++i)
    cin >> weights[i];

  iota(fu_parent, fu_parent + N + 1, 0);

  for (int i = 0; i < N - 1; ++i) {
    int u, v, c;
    cin >> u >> v >> c;
    adj[u].emplace_back(v, c);
    adj[v].emplace_back(u, c);
  }

  fill_n(is_avail, N + 1, true);
  fill_n(layer_idx, N + 1, (int)1e9);
  Centroids(1, 0);

  int num_edges = 0;
  long long total_cost = 0;
  iota(groups, groups + N + 1, 0);

  while (num_edges < N - 1) {
    CentroidGroups();

    vector<int> order(N);
    vector<bool> had_out_edge(N + 1);
    iota(ALL(order), 1);
    sort(ALL(order), [&](int lhs, int rhs) {
      return vert_ptr_dist[lhs] < vert_ptr_dist[rhs];
    });

    for (int i : order) {
      if (had_out_edge[groups[i]]) {
        continue;
      }
      had_out_edge[groups[i]] = true;
      debug(i, vert_ptr[i], vert_ptr_dist[i]);
      assert(vert_ptr_dist[i] < 1e17);
      if (FuUnion(i, vert_ptr[i])) {
        total_cost += vert_ptr_dist[i];
        ++num_edges;
        debug("TAKING EDGE", i, vert_ptr[i], vert_ptr_dist[i],
              "!@#@!@!#@!##@!");
      }
    }

    for (int i = 1; i <= N; ++i) {
      groups[i] = FuFind(i);
      debug(i, groups[i]);
    }
  }

  cout << total_cost << "\\n";
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>

#pragma GCC optimize("O3")

using namespace std;
#define PB push_back
#define MP make_pair
#define LL long long
// #define int LL
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define RE(i, n) FOR(i, 1, n)
#define REP(i, n) FOR(i, 0, (int)(n)-1)
#define R(i, n) REP(i, n)
#define VI vector<int>
#define PII pair<int, int>
#define LD long double
#define FI first
#define SE second
#define st FI
#define nd SE
#define ALL(x) (x).begin(), (x).end()
#define SZ(x) ((int)(x).size())

#define unordered_map __fast_unordered_map
template <class Key, class Value, class Hash = std::hash<Key>>
using unordered_map = __gnu_pbds::gp_hash_table<Key, Value, Hash>;

template <class C> void mini(C &a4, C b4) { a4 = min(a4, b4); }
template <class C> void maxi(C &a4, C b4) { a4 = max(a4, b4); }

template <class TH> void _dbg(const char *sdbg, TH h) {
  cerr << sdbg << '=' << h << endl;
}
template <class TH, class... TA> void _dbg(const char *sdbg, TH h, TA... a) {
  while (*sdbg != ',')
    cerr << *sdbg++;
  cerr << '=' << h << ',';
  _dbg(sdbg + 1, a...);
}

template <class T> ostream &operator<<(ostream &os, vector<T> V) {
  os << "[";
  for (auto vv : V)
    os << vv << ",";
  return os << "]";
}
template <class L, class R> ostream &operator<<(ostream &os, pair<L, R> P) {
  return os << "(" << P.st << "," << P.nd << ")";
}
template <class L, class R, class X>
ostream &operator<<(ostream &os, tuple<L, R, X> P) {
  return os << "(" << get<0>(P) << "," << get<1>(P) << "," << get<2>(P) << ")";
}

#ifdef LOCAL
#define debug(...) _dbg(#__VA_ARGS__, __VA_ARGS__)
#else
#define debug(...) (__VA_ARGS__)
#define cerr                                                                   \\
  if (0)                                                                       \\
  cout
#endif

const int MaxN = 2e5 + 100;

int groups[MaxN];
int centroid_subtree[MaxN];

struct BestInfo {
  LL dist;
  // int subtree_idx;
  int vert;

  bool operator<(const BestInfo &other) const { return dist < other.dist; }

  bool SameType(const BestInfo &other) const {
    return groups[vert] == groups[other.vert];
  }

  friend ostream &operator<<(ostream &os, const BestInfo &bi) {
    return os << "(" << bi.dist << "," << bi.vert << ")";
  }
};

vector<PII> adj[MaxN];
int weights[MaxN];
int N;
bool is_avail[MaxN];
int tree_size[MaxN];
int layer_idx[MaxN];
vector<int> conn_verts;
vector<BestInfo> best_dists[MaxN];

void DfsSizes(int v, int p) {
  tree_size[v] = 1;
  conn_verts.push_back(v);
  for (auto &E : adj[v]) {
    const int s = E.st;
    if (s == p) {
      continue;
    }
    if (!is_avail[s]) {
      continue;
    }
    DfsSizes(s, v);
    tree_size[v] += tree_size[s];
  }
}

void Centroids(int v, int layer) {
  conn_verts.clear();
  DfsSizes(v, v);

  int centroid = v;
  for (int s : conn_verts)
    if (2 * tree_size[s] >= tree_size[v] && tree_size[s] < tree_size[centroid])
      centroid = s;

  is_avail[centroid] = false;
  layer_idx[centroid] = layer;
  debug(centroid, layer);

  for (auto &E : adj[centroid])
    if (is_avail[E.st])
      Centroids(E.st, layer + 1);
}

const int kMaxBests = 4;
int subtree_idx;
int centr_cur;
LL dists_from_centroid[MaxN];
int vert_ptr[MaxN];
LL vert_ptr_dist[MaxN];

void DfsDists(int v, int p, LL cur_dist, vector<BestInfo> &bests) {
  debug(v, p, centr_cur);
  dists_from_centroid[v] = cur_dist;
  conn_verts.push_back(v);
  centroid_subtree[v] = subtree_idx;

  if (v != p) {
    bool changed = false;
    BestInfo my_info{cur_dist + weights[v], v};

    for (auto &x : bests) {
      if (x.SameType(my_info)) {
        changed = true;
        if (x.dist > my_info.dist)
          x = my_info;
        break;
      }
    }

    if (!changed) {
      bests.push_back(my_info);
      swap(*min_element(ALL(bests)), bests[0]);
      int lptr = 1;
      int seen[2] = {-1, 0};
      for (int i = 1; i < SZ(bests); ++i) {
        const bool same_group = groups[bests[i].vert] == groups[bests[0].vert];
        if (!same_group && bests[i].dist <= bests[1].dist) {
          bests[1] = bests[i];
          ++lptr;
        }
      }
      bests.resize(lptr);
    } else {
      // sort(ALL(bests));
    }

    if (centr_cur == 2)
      debug(v, bests, cur_dist, weights[v]);
  }

  for (auto &E : adj[v]) {
    const int s = E.st;
    if (s == p || layer_idx[s] < layer_idx[centr_cur]) {
      continue;
    }
    if (v == p) {
      subtree_idx = s;
    }
    DfsDists(s, v, cur_dist + E.nd, bests);
  }
}

void CentroidGroups() {
  fill_n(vert_ptr_dist, N + 1, (LL)1e18);

  for (int centroid = 1; centroid <= N; ++centroid) {
    best_dists[centroid].clear();
    subtree_idx = -1;
    centr_cur = centroid;
    conn_verts.clear();
    DfsDists(centroid, centroid, 0, best_dists[centroid]);
    debug(centroid, best_dists[centroid]);
    debug(conn_verts);

    for (int v : conn_verts) {
      for (auto &best : best_dists[centroid]) {
        if (groups[best.vert] == groups[v])
          continue;

        const LL new_dist = best.dist + dists_from_centroid[v] + weights[v];
        debug(centroid, v, dists_from_centroid[v], best.vert, new_dist);
        if (new_dist < vert_ptr_dist[v]) {
          vert_ptr[v] = best.vert;
          vert_ptr_dist[v] = new_dist;
        }
      }

      if (groups[centroid] != groups[v]) {
        const LL centr_dist =
            dists_from_centroid[v] + weights[v] + weights[centroid];
        if (centr_dist < vert_ptr_dist[v]) {
          vert_ptr[v] = centroid;
          vert_ptr_dist[v] = centr_dist;
        }
        if (centr_dist < vert_ptr_dist[centroid]) {
          vert_ptr[centroid] = v;
          vert_ptr_dist[centroid] = centr_dist;
        }
      }
    }
  }
}

int fu_parent[MaxN];

int FuFind(int v) {
  if (fu_parent[v] == v)
    return v;
  return fu_parent[v] = FuFind(fu_parent[v]);
}

bool FuUnion(int a, int b) {
  a = FuFind(a);
  b = FuFind(b);
  if (a != b) {
    fu_parent[a] = b;
    return true;
  }
  return false;
}

int32_t main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout << fixed << setprecision(11);
  cerr << fixed << setprecision(6);

  cin >> N;
  for (int i = 1; i <= N; ++i)
    cin >> weights[i];

  iota(fu_parent, fu_parent + N + 1, 0);

  for (int i = 0; i < N - 1; ++i) {
    int u, v, c;
    cin >> u >> v >> c;
    adj[u].emplace_back(v, c);
    adj[v].emplace_back(u, c);
  }

  fill_n(is_avail, N + 1, true);
  fill_n(layer_idx, N + 1, (int)1e9);
  Centroids(1, 0);

  int num_edges = 0;
  long long total_cost = 0;
  iota(groups, groups + N + 1, 0);

  while (num_edges < N - 1) {
    CentroidGroups();

    vector<int> order(N);
    vector<bool> had_out_edge(N + 1);
    iota(ALL(order), 1);
    sort(ALL(order), [&](int lhs, int rhs) {
      return vert_ptr_dist[lhs] < vert_ptr_dist[rhs];
    });

    for (int i : order) {
      if (had_out_edge[groups[i]]) {
        continue;
      }
      had_out_edge[groups[i]] = true;
      debug(i, vert_ptr[i], vert_ptr_dist[i]);
      assert(vert_ptr_dist[i] < 1e17);
      if (FuUnion(i, vert_ptr[i])) {
        total_cost += vert_ptr_dist[i];
        ++num_edges;
        debug("TAKING EDGE", i, vert_ptr[i], vert_ptr_dist[i],
              "!@#@!@!#@!##@!");
      }
    }

    for (int i = 1; i <= N; ++i) {
      groups[i] = FuFind(i);
      debug(i, groups[i]);
    }
  }

  cout << total_cost << "\\n";
}