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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>

#define pb push_back
#define mp make_pair

using namespace std;

template <typename T> inline bool upmin(T &x, T y) {
  return y < x ? x = y, 1 : 0;
}

typedef unsigned int u32;
typedef long long LL;
typedef unsigned long long ULL;
typedef long double lod;
typedef pair<int, int> PR;
typedef vector<int> VI;

const lod pi = acos(-1);
const int oo = 1 << 30;
const LL OO = 1e18;

const int N = 2e5 + 100;

int gi() {
  int w = 0;
  bool q = 1;
  char c = getchar();
  while ((c < '0' || c > '9') && c != '-')
    c = getchar();
  if (c == '-')
    q = 0, c = getchar();
  while (c >= '0' && c <= '9')
    w = w * 10 + c - '0', c = getchar();
  return q ? w : -w;
}

int val[N], fa[N];
int head[N], nxt[N * 2], to[N * 2], w[N * 2];
inline int find(int x) { return fa[x] == x ? x : fa[x] = find(fa[x]); }
#define fi first
#define se second
struct DT {
  pair<LL, int> fi, se;
} f[N];
inline DT operator+(const DT &a, const int &b) {
  return (DT){mp(a.fi.fi + b, a.fi.se), mp(a.se.fi + b, a.se.se)};
}
inline DT operator+(const DT &a, const DT &b) {
  DT ans;
  if (a.fi.se == b.fi.se)
    ans.fi = min(a.fi, b.fi), ans.se = min(a.se, b.se);
  else if (a.fi.fi < b.fi.fi)
    ans = a, upmin(ans.se, b.fi);
  else
    ans = b, upmin(ans.se, a.fi);
  return ans;
}
inline void dfs1(int k, int fa) {
  f[k] = (DT){mp(val[k], find(k)), mp(OO, 0)};
  for (int i = head[k]; i; i = nxt[i])
    if (to[i] != fa)
      dfs1(to[i], k), f[k] = f[k] + (f[to[i]] + w[i]);
}
inline void dfs2(int k, int fa) {
  for (int i = head[k]; i; i = nxt[i])
    if (to[i] != fa) {
      f[to[i]] = f[to[i]] + (f[k] + w[i]);
      dfs2(to[i], k);
    }
}
int main() {
  int n = gi(), i, a, b, C = n, tot = 0;
  pair<LL, int> p;
  LL ans = 0;
  for (i = 1; i <= n; i++)
    val[i] = gi(), fa[i] = i;
  for (i = 1; i < n; i++) {
    a = gi(), b = gi();
    to[++tot] = b, nxt[tot] = head[a], head[a] = tot;
    to[++tot] = a, nxt[tot] = head[b], head[b] = tot;
    w[tot] = w[tot - 1] = gi();
  }
  while (C > 1) {
    dfs1(1, 0);
    dfs2(1, 0);
    for (i = 1; i <= n; i++)
      f[i].fi.fi += val[i], f[i].se.fi += val[i];
    for (i = 1; i <= n; i++)
      if (fa[i] != i)
        f[find(i)] = f[find(i)] + f[i];
    for (i = 1; i <= n; i++)
      if (fa[i] == i) {
        p = f[i].fi.se == i ? f[i].se : f[i].fi;
        if (find(i) != find(p.se)) {
          ans += p.fi;
          fa[i] = p.se;
          C--;
        }
      }
  }
  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>

#define pb push_back
#define mp make_pair

using namespace std;

template <typename T> inline bool upmin(T &x, T y) {
  return y < x ? x = y, 1 : 0;
}

typedef unsigned int u32;
typedef long long LL;
typedef unsigned long long ULL;
typedef long double lod;
typedef pair<int, int> PR;
typedef vector<int> VI;

const lod pi = acos(-1);
const int oo = 1 << 30;
const LL OO = 1e18;

const int N = 2e5 + 100;

int gi() {
  int w = 0;
  bool q = 1;
  char c = getchar();
  while ((c < '0' || c > '9') && c != '-')
    c = getchar();
  if (c == '-')
    q = 0, c = getchar();
  while (c >= '0' && c <= '9')
    w = w * 10 + c - '0', c = getchar();
  return q ? w : -w;
}

int val[N], fa[N];
int head[N], nxt[N * 2], to[N * 2], w[N * 2];
inline int find(int x) { return fa[x] == x ? x : fa[x] = find(fa[x]); }
#define fi first
#define se second
struct DT {
  pair<LL, int> fi, se;
} f[N];
inline DT operator+(const DT &a, const int &b) {
  return (DT){mp(a.fi.fi + b, a.fi.se), mp(a.se.fi + b, a.se.se)};
}
inline DT operator+(const DT &a, const DT &b) {
  DT ans;
  if (a.fi.se == b.fi.se)
    ans.fi = min(a.fi, b.fi), ans.se = min(a.se, b.se);
  else if (a.fi.fi < b.fi.fi)
    ans = a, upmin(ans.se, b.fi);
  else
    ans = b, upmin(ans.se, a.fi);
  return ans;
}
inline void dfs1(int k, int fa) {
  f[k] = (DT){mp(val[k], find(k)), mp(OO, 0)};
  for (int i = head[k]; i; i = nxt[i])
    if (to[i] != fa)
      dfs1(to[i], k), f[k] = f[k] + (f[to[i]] + w[i]);
}
inline void dfs2(int k, int fa) {
  for (int i = head[k]; i; i = nxt[i])
    if (to[i] != fa) {
      f[to[i]] = f[to[i]] + (f[k] + w[i]);
      dfs2(to[i], k);
    }
}
int main() {
#ifndef ONLINE_JUDGE
  freopen("J.in", "r", stdin);
  freopen("J.out", "w", stdout);
#endif
  int n = gi(), i, a, b, C = n, tot = 0;
  pair<LL, int> p;
  LL ans = 0;
  for (i = 1; i <= n; i++)
    val[i] = gi(), fa[i] = i;
  for (i = 1; i < n; i++) {
    a = gi(), b = gi();
    to[++tot] = b, nxt[tot] = head[a], head[a] = tot;
    to[++tot] = a, nxt[tot] = head[b], head[b] = tot;
    w[tot] = w[tot - 1] = gi();
  }
  while (C > 1) {
    dfs1(1, 0);
    dfs2(1, 0);
    for (i = 1; i <= n; i++)
      f[i].fi.fi += val[i], f[i].se.fi += val[i];
    for (i = 1; i <= n; i++)
      if (fa[i] != i)
        f[find(i)] = f[find(i)] + f[i];
    for (i = 1; i <= n; i++)
      if (fa[i] == i) {
        p = f[i].fi.se == i ? f[i].se : f[i].fi;
        if (find(i) != find(p.se)) {
          ans += p.fi;
          fa[i] = p.se;
          C--;
        }
      }
  }
  cout << ans << endl;
  return 0;
}