Problem p03532 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
template <class T> inline void init(T &x) {
  x = 0;
  char ch = getchar();
  bool t = 0;
  for (; ch > '9' || ch < '0'; ch = getchar())
    if (ch == '-')
      t = 1;
  for (; ch >= '0' && ch <= '9'; ch = getchar())
    x = (x << 1) + (x << 3) + (ch - 48);
  if (t)
    x = -x;
  return;
}
const int N = 2e5 + 10;
const int M = 1e7;
typedef long long ll;
struct edge {
  int u, v;
  ll val;
  inline bool operator<(const edge B) const { return val < B.val; }
} E[M];
int cur = 0;
int fa[N];
ll w[N];
int n, m;
struct Edge {
  int to, next;
  ll w;
} a[N << 1];
int head[N], cnt = 0;
inline void add(int x, int y, ll w) {
  a[++cnt] = (Edge){y, head[x], w};
  head[x] = cnt;
}
int size[N], f[N];
bool vis[N];
int RT, SZ;
void Find(int u, int fa) {
  size[u] = 1, f[u] = 0;
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (v == fa || vis[v])
      continue;
    Find(v, u);
    f[u] = max(f[u], size[v]);
    size[u] += size[v];
  }
  f[u] = max(f[u], SZ - size[u]);
  if ((!RT) || f[RT] > f[u])
    RT = u;
  return;
}
typedef pair<ll, int> PA;
PA Dis[N];
int top = 0;
PA Mi;
void Dfs(int u, int fa, ll dis) {
  Dis[++top] = PA(dis + w[u], u);
  if (Mi > Dis[top])
    Mi = Dis[top];
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (vis[v] || v == fa)
      continue;
    Dfs(v, u, dis + a[i].w);
  }
  return;
}
int dep = 0, mx = 0;
void Divide(int u) {
  vis[u] = 1;
  top = 0;
  Mi = PA(w[u], u);
  ++dep;
  mx = max(dep, mx);
  int L = 0, R = 0;
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (vis[v])
      continue;
    L = top;
    PA now = Mi;
    Dfs(v, u, a[i].w);
    if (now != Mi)
      R = top;
  }
  if (!R) {
    int U = Mi.second;
    ll val = Mi.first;
    for (int i = 1; i <= top; ++i) {
      int V = Dis[i].second;
      E[++cur] = (edge){U, V, val + Dis[i].first};
    }
  } else {
    E[++cur] = (edge){u, Mi.second, Mi.first + w[u]};
    int U = Mi.second;
    ll val = Mi.first;
    for (int i = 1; i <= top; ++i) {
      if (i > L && i <= R)
        continue;
      int V = Dis[i].second;
      E[++cur] = (edge){U, V, val + Dis[i].first};
    }
  }
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (vis[v])
      continue;
    RT = 0;
    SZ = size[v];
    Find(v, u);
    Divide(RT);
  }
  --dep;
  return;
}
int find(int x) { return fa[x] == x ? x : fa[x] = find(fa[x]); }
int main() {
  freopen("data.in", "r", stdin);
  init(n);
  for (int i = 1; i <= n; ++i)
    init(w[i]);
  int u, v;
  ll val;
  for (int i = 1; i < n; ++i) {
    init(u), init(v), init(val);
    add(u, v, val), add(v, u, val);
  }
  SZ = n;
  Find(1, 0);
  Divide(RT);
  sort(E + 1, E + 1 + cur);
  for (int i = 1; i <= n; ++i)
    fa[i] = i;
  ll ans = 0;
  for (int i = 1; i <= cur; ++i) {
    int u = E[i].u, v = E[i].v;
    ll Val = E[i].val;
    int fx = find(u), fy = find(v);
    if (fx == fy)
      continue;
    fa[fx] = fy;
    ans += Val;
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
template <class T> inline void init(T &x) {
  x = 0;
  char ch = getchar();
  bool t = 0;
  for (; ch > '9' || ch < '0'; ch = getchar())
    if (ch == '-')
      t = 1;
  for (; ch >= '0' && ch <= '9'; ch = getchar())
    x = (x << 1) + (x << 3) + (ch - 48);
  if (t)
    x = -x;
  return;
}
const int N = 2e5 + 10;
const int M = 1e7;
typedef long long ll;
struct edge {
  int u, v;
  ll val;
  inline bool operator<(const edge B) const { return val < B.val; }
} E[M];
int cur = 0;
int fa[N];
ll w[N];
int n, m;
struct Edge {
  int to, next;
  ll w;
} a[N << 1];
int head[N], cnt = 0;
inline void add(int x, int y, ll w) {
  a[++cnt] = (Edge){y, head[x], w};
  head[x] = cnt;
}
int size[N], f[N];
bool vis[N];
int RT, SZ;
void Find(int u, int fa) {
  size[u] = 1, f[u] = 0;
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (v == fa || vis[v])
      continue;
    Find(v, u);
    f[u] = max(f[u], size[v]);
    size[u] += size[v];
  }
  f[u] = max(f[u], SZ - size[u]);
  if ((!RT) || f[RT] > f[u])
    RT = u;
  return;
}
typedef pair<ll, int> PA;
PA Dis[N];
int top = 0;
PA Mi;
void Dfs(int u, int fa, ll dis) {
  Dis[++top] = PA(dis + w[u], u);
  if (Mi > Dis[top])
    Mi = Dis[top];
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (vis[v] || v == fa)
      continue;
    Dfs(v, u, dis + a[i].w);
  }
  return;
}
int dep = 0, mx = 0;
void Divide(int u) {
  vis[u] = 1;
  top = 0;
  Mi = PA(w[u], u);
  ++dep;
  mx = max(dep, mx);
  int L = 0, R = 0;
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (vis[v])
      continue;
    L = top;
    PA now = Mi;
    Dfs(v, u, a[i].w);
    if (now != Mi)
      R = top;
  }
  if (!R) {
    int U = Mi.second;
    ll val = Mi.first;
    for (int i = 1; i <= top; ++i) {
      int V = Dis[i].second;
      E[++cur] = (edge){U, V, val + Dis[i].first};
    }
  } else {
    E[++cur] = (edge){u, Mi.second, Mi.first + w[u]};
    int U = Mi.second;
    ll val = Mi.first;
    for (int i = 1; i <= top; ++i) {
      if (i > L && i <= R)
        continue;
      int V = Dis[i].second;
      E[++cur] = (edge){U, V, val + Dis[i].first};
    }
  }
  for (int v, i = head[u]; i; i = a[i].next) {
    v = a[i].to;
    if (vis[v])
      continue;
    RT = 0;
    SZ = size[v];
    Find(v, u);
    Divide(RT);
  }
  --dep;
  return;
}
int find(int x) { return fa[x] == x ? x : fa[x] = find(fa[x]); }
int main() {
  init(n);
  for (int i = 1; i <= n; ++i)
    init(w[i]);
  int u, v;
  ll val;
  for (int i = 1; i < n; ++i) {
    init(u), init(v), init(val);
    add(u, v, val), add(v, u, val);
  }
  SZ = n;
  Find(1, 0);
  Divide(RT);
  sort(E + 1, E + 1 + cur);
  for (int i = 1; i <= n; ++i)
    fa[i] = i;
  ll ans = 0;
  for (int i = 1; i <= cur; ++i) {
    int u = E[i].u, v = E[i].v;
    ll Val = E[i].val;
    int fx = find(u), fy = find(v);
    if (fx == fy)
      continue;
    fa[fx] = fy;
    ans += Val;
  }
  cout << ans << endl;
  return 0;
}