Problem p03542 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>

#define pb push_back
#define mp make_pair
#define taskname "A"

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int, int> ii;

const int maxn = 2e5 + 5;
const int logn = log2(maxn) + 2;
const int mod = 1e9 + 7;
int lab[maxn];
int FindLab(int u) { return lab[u] < 0 ? u : lab[u] = FindLab(lab[u]); }
int n, cost[maxn];
struct node {
  pair<ll, int> a, b;
  node() { a = b = mp(1e18, -1); };
  void reset() { a = b = mp(1e18, -1); }
  void add(pair<ll, int> x, ll y = 0) {
    x.first += y;
    if (a.second == x.second)
      a = min(a, x);
    else if (b.second == x.second)
      b = min(b, x);
    else if (a > x)
      b = a, a = x;
    else if (b > x)
      b = x;
    if (b < a)
      swap(a, b);
    assert(b.second != a.second || a.second == -1);
  }
} dp[maxn];
pair<ll, int> bezt[maxn];
vector<ii> adj[maxn];

node dfs(int u, int par) {
  node a;
  //    cout << u << endl;
  for (auto c : adj[u]) {
    if (c.first == par)
      continue;
    auto tmp = dfs(c.first, u);
    a.add(tmp.a, c.second);
    a.add(tmp.b, c.second);
  }
  if (a.a.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.a.first + cost[u], a.a.second));
  if (a.b.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.b.first + cost[u], a.b.second));
  a.add(mp(cost[u], FindLab(u)));
  for (auto c : adj[u]) {
    if (c.first == par)
      continue;
    dp[c.first].add(a.a);
    dp[c.first].add(a.b);
  }
  return a;
}

void dfs1(int u, int par, node a) {
  if (a.a.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.a.first + cost[u], a.a.second));
  if (a.b.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.b.first + cost[u], a.b.second));
  for (auto c : adj[u]) {
    if (c.first == par)
      continue;
    node b;
    b.add(a.a, c.second);
    b.add(a.b, c.second);
    b.add(dp[c.first].a, c.second);
    b.add(dp[c.first].b, c.second);
    dfs1(c.first, u, b);
  }
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  if (fopen(taskname ".INP", "r")) {
    freopen(taskname ".INP", "r", stdin);
    freopen(taskname ".OUT", "w", stdout);
  }
  fill_n(lab, maxn, -1);
  cin >> n;
  for (int i = 1; i <= n; ++i)
    cin >> cost[i];
  for (int i = 1; i < n; ++i) {
    int u, v, c;
    cin >> u >> v >> c;
    adj[u].pb(mp(v, c));
    adj[v].pb(mp(u, c));
  }
  ll res = 0;
  int cnt = n - 1;
  while (cnt > 0) {
    for (int i = 1; i <= n; ++i)
      dp[i].reset(), bezt[i] = mp(1e18, -1);
    dfs(1, 0);
    //        return 0;
    node tmp;
    tmp.reset();
    dfs1(1, 0, tmp);
    vector<pair<ll, ii>> edge;
    for (int i = 1; i <= n; ++i) {
      if (lab[i] < 0) {
        if (bezt[i].first < 1e18) {
          edge.pb(mp(bezt[i].first, mp(i, bezt[i].second)));
          int u = i, v = bezt[i].second;
          if (u == v)
            continue;
          res += bezt[i].first;
          if (lab[u] > lab[v])
            swap(u, v);
          lab[u] += lab[v];
          lab[v] = u;
          --cnt;
        }
      }
    }
  }
  cout << res;
}


Fixed submission:
#include <bits/stdc++.h>

#define pb push_back
#define mp make_pair
#define taskname "A"

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int, int> ii;

const int maxn = 2e5 + 5;
const int logn = log2(maxn) + 2;
const int mod = 1e9 + 7;
int lab[maxn];
int FindLab(int u) { return lab[u] < 0 ? u : lab[u] = FindLab(lab[u]); }
int n, cost[maxn];
struct node {
  pair<ll, int> a, b;
  node() { a = b = mp(1e18, -1); };
  void reset() { a = b = mp(1e18, -1); }
  void add(pair<ll, int> x, ll y = 0) {
    x.first += y;
    if (a.second == x.second)
      a = min(a, x);
    else if (b.second == x.second)
      b = min(b, x);
    else if (a > x)
      b = a, a = x;
    else if (b > x)
      b = x;
    if (b < a)
      swap(a, b);
    assert(b.second != a.second || a.second == -1);
  }
} dp[maxn];
pair<ll, int> bezt[maxn];
vector<ii> adj[maxn];

node dfs(int u, int par) {
  node a;
  //    cout << u << endl;
  for (auto c : adj[u]) {
    if (c.first == par)
      continue;
    auto tmp = dfs(c.first, u);
    a.add(tmp.a, c.second);
    a.add(tmp.b, c.second);
  }
  if (a.a.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.a.first + cost[u], a.a.second));
  if (a.b.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.b.first + cost[u], a.b.second));
  a.add(mp(cost[u], FindLab(u)));
  for (auto c : adj[u]) {
    if (c.first == par)
      continue;
    dp[c.first].add(a.a);
    dp[c.first].add(a.b);
  }
  return a;
}

void dfs1(int u, int par, node a) {
  if (a.a.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.a.first + cost[u], a.a.second));
  if (a.b.second != FindLab(u))
    bezt[FindLab(u)] =
        min(bezt[FindLab(u)], mp(a.b.first + cost[u], a.b.second));
  for (auto c : adj[u]) {
    if (c.first == par)
      continue;
    node b;
    b.add(a.a, c.second);
    b.add(a.b, c.second);
    b.add(dp[c.first].a, c.second);
    b.add(dp[c.first].b, c.second);
    dfs1(c.first, u, b);
  }
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  if (fopen(taskname ".INP", "r")) {
    freopen(taskname ".INP", "r", stdin);
    freopen(taskname ".OUT", "w", stdout);
  }
  fill_n(lab, maxn, -1);
  cin >> n;
  for (int i = 1; i <= n; ++i)
    cin >> cost[i];
  for (int i = 1; i < n; ++i) {
    int u, v, c;
    cin >> u >> v >> c;
    adj[u].pb(mp(v, c));
    adj[v].pb(mp(u, c));
  }
  ll res = 0;
  int cnt = n - 1;
  while (cnt > 0) {
    for (int i = 1; i <= n; ++i)
      dp[i].reset(), bezt[i] = mp(1e18, -1);
    dfs(1, 0);
    //        return 0;
    node tmp;
    tmp.reset();
    dfs1(1, 0, tmp);
    vector<pair<ll, ii>> edge;
    for (int i = 1; i <= n; ++i) {
      if (lab[i] < 0) {
        if (bezt[i].first < 1e18) {
          edge.pb(mp(bezt[i].first, mp(i, bezt[i].second)));
          int u = FindLab(i), v = FindLab(bezt[i].second);
          if (u == v)
            continue;
          res += bezt[i].first;
          if (lab[u] > lab[v])
            swap(u, v);
          lab[u] += lab[v];
          lab[v] = u;
          --cnt;
        }
      }
    }
  }
  cout << res;
}