Problem p03532 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <deque>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
#define mk make_pair
#define pb push_back
#define fi first
#define se second
#define REP(i, x, y) for (int i = (int)x; i <= (int)y; i++)
#define FOR(i, x, y) for (int i = (int)x; i < (int)y; i++)
#define PER(i, x, y) for (int i = (int)x; i >= (int)y; i--)
#define trace(x) cerr << #x << " " << x << endl;
#define dprintf(...) fprintf(stderr, __VA__ARGS__)
#define dln() fprintf(stderr, "\n")
using namespace std;
typedef double db;
typedef long long LL;
typedef pair<int, int> PII;
typedef vector<int> VI;
typedef vector<VI> VII;
typedef vector<PII> VPI;
typedef vector<long long> VL;
const int N = 200005;
const int P = 1e9 + 7;
const int inf = 1e9;
const LL Inf = 1e18;

inline int IN() {
  char ch = getchar();
  int x = 0, f = 0;
  while (ch < '0' || ch > '9')
    ch = getchar(), f = (ch == '-');
  while (ch >= '0' && ch <= '9') {
    x = (x << 1) + (x << 3) + ch - 48;
    ch = getchar();
  }
  return f ? (-x) : x;
}

inline int Pow(int x, int y, int p) {
  int an = 1;
  for (; y; y >>= 1, x = (LL)x * x % p)
    if (y & 1)
      an = (LL)an * x % p;
  return an;
}

void renew(int &x, int y) {
  x += y;
  if (x < 0)
    x += P;
  else if (x >= P)
    x -= P;
}

void setIO(string a) {
#ifndef LOCAL
  freopen((a + ".in").c_str(), "r", stdin);
  freopen((a + ".out").c_str(), "w", stdout);
#endif
}

template <typename T> inline void chkmin(T &a, const T &b) {
  if (a > b)
    a = b;
}
template <typename T> inline void chkmax(T &a, const T &b) {
  if (a < b)
    a = b;
}

LL X[N];
int cnt, n, col[N];
struct edge {
  int v;
  LL w;
  int pre;
} e[N * 2];
int dex, adj[N];

struct info {
  LL d1;
  int c1, go1;
  LL d2;
  int c2, go2;
  info operator+(const info &a) const {
    info ans = *this;
    if (a.d1 < ans.d1) {
      if (ans.c1 == a.c1) {
        ans.go1 = a.go1;
        ans.d1 = a.d1;
      } else {
        ans.d2 = ans.d1;
        ans.c2 = ans.c1;
        ans.go2 = ans.go1;
        ans.d1 = a.d1;
        ans.c1 = a.c1;
        ans.go1 = a.go1;
      }
    }
    if (a.d1 < ans.d2 && a.c1 != ans.c1) {
      ans.d2 = a.d1;
      ans.c2 = a.c1;
      ans.go2 = a.go1;
    }
    if (a.d2 < ans.d2 && a.c2 != ans.c1) {
      ans.d2 = a.d2;
      ans.c2 = a.c2;
      ans.go2 = a.go2;
    }
    return ans;
  }
} dp1[N], dp2[N], seq1[N], seq2[N], tmp[N];
int L[N], R[N], tot, go[N], ID[N], fa[N];

void dfs(int x, int p = 0) {
  int fir = 1;
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    dfs(v, x);
  }
  L[x] = tot + 1;
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    R[x] = ++tot;
    go[tot] = v;
    ID[v] = tot;
    tmp[tot] = dp1[v];
    tmp[tot].d1 += e[i].w;
    tmp[tot].d2 += e[i].w;
  }
  seq1[L[x]] = tmp[L[x]];
  seq2[R[x]] = tmp[R[x]];
  REP(i, L[x] + 1, R[x]) seq1[i] = seq1[i - 1] + tmp[i];
  PER(i, R[x] - 1, L[x]) seq2[i] = seq2[i + 1] + tmp[i];
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    if (fir) {
      info ww = dp1[v];
      ww.d1 += e[i].w;
      ww.d2 += e[i].w;
      dp1[x] = dp1[x] + ww;
      fir = 0;
    } else {
      info ww = dp1[v];
      ww.d1 += e[i].w;
      ww.d2 += e[i].w;
      dp1[x] = dp1[x] + ww;
    }
  }
}

int nxt[N];
LL ww[N];

LL dd[10005];
int dst = 0;
void _dfs(int x, int y) {
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == y)
      continue;
    dd[v] = dd[x] + e[i].w;
    _dfs(v, x);
  }
}
int chk(int u, int v, LL w) {
  REP(i, 1, n) dd[i] = 0;
  _dfs(u, 0);
  assert(dd[v] == w - X[v]);
}

void dfs2(int x, int p = 0) {
  dp1[x] = dp1[x] + dp2[x];
  info uu = dp1[x];
  assert(dp1[x].c1 != dp1[x].c2);
  assert(uu.d1 <= uu.d2);
  chk(x, dp1[x].go1, dp1[x].d1);
  chk(x, dp1[x].go2, dp1[x].d2);
  if (uu.c1 == col[x]) {
    if (uu.d2 + X[x] < ww[col[x]]) {
      ww[col[x]] = uu.d2 + X[x];
      nxt[col[x]] = uu.c2;
    }
  } else {
    if (uu.d1 + X[x] < ww[col[x]]) {
      ww[col[x]] = uu.d1 + X[x];
      nxt[col[x]] = uu.c1;
    }
  }
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    info now = dp2[x];
    if (ID[v] != L[x])
      now = now + seq1[ID[v] - 1];
    if (ID[v] != R[x])
      now = now + seq2[ID[v] + 1];
    now.d1 += e[i].w;
    now.d2 += e[i].w;
    dp2[v] = dp2[v] + now;
    dfs2(v, x);
  }
}

int get(int x) {
  if (fa[x] == x)
    return x;
  return fa[x] = get(fa[x]);
}

int app[N];

int main() {
#ifdef LOCAL
  freopen("1.in", "r", stdin);
//	freopen("1.out", "w", stdout);
#endif

  scanf("%d", &n);
  REP(i, 1, n) scanf("%lld", X + i);
  REP(i, 1, n - 1) {
    int x, y, z;
    scanf("%d%d%d", &x, &y, &z);
    e[++dex] = (edge){y, z, adj[x]};
    adj[x] = dex;
    e[++dex] = (edge){x, z, adj[y]};
    adj[y] = dex;
  }

  REP(i, 1, n) col[i] = i;
  cnt = n;
  LL res = 0;
  while (cnt > 1) {
    REP(i, 1, cnt) ww[i] = 1e18;
    REP(i, 1, n) {
      dp1[i] = (info){X[i], col[i], i, (LL)1e18, 0, 0};
      dp2[i] = (info){X[i], col[i], i, (LL)1e18, 0, 0};
    }
    tot = 0;
    dfs(1, 0);
    dfs2(1, 0);
    REP(i, 1, cnt) fa[i] = i;
    REP(i, 1, cnt) {
      int u = get(i), v = get(nxt[i]);
      if (u == v)
        continue;
      res += ww[i];
      fa[u] = v;
    }
    REP(i, 1, n) col[i] = get(col[i]);
    REP(i, 1, cnt) app[i] = 0;
    cnt = 0;
    REP(i, 1, n) if (app[col[i]]) col[i] = app[col[i]];
    else col[i] = app[col[i]] = ++cnt;
  }

  cout << res << endl;

  return 0;
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <deque>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
#define mk make_pair
#define pb push_back
#define fi first
#define se second
#define REP(i, x, y) for (int i = (int)x; i <= (int)y; i++)
#define FOR(i, x, y) for (int i = (int)x; i < (int)y; i++)
#define PER(i, x, y) for (int i = (int)x; i >= (int)y; i--)
#define trace(x) cerr << #x << " " << x << endl;
#define dprintf(...) fprintf(stderr, __VA__ARGS__)
#define dln() fprintf(stderr, "\n")
using namespace std;
typedef double db;
typedef long long LL;
typedef pair<int, int> PII;
typedef vector<int> VI;
typedef vector<VI> VII;
typedef vector<PII> VPI;
typedef vector<long long> VL;
const int N = 200005;
const int P = 1e9 + 7;
const int inf = 1e9;
const LL Inf = 1e18;

inline int IN() {
  char ch = getchar();
  int x = 0, f = 0;
  while (ch < '0' || ch > '9')
    ch = getchar(), f = (ch == '-');
  while (ch >= '0' && ch <= '9') {
    x = (x << 1) + (x << 3) + ch - 48;
    ch = getchar();
  }
  return f ? (-x) : x;
}

inline int Pow(int x, int y, int p) {
  int an = 1;
  for (; y; y >>= 1, x = (LL)x * x % p)
    if (y & 1)
      an = (LL)an * x % p;
  return an;
}

void renew(int &x, int y) {
  x += y;
  if (x < 0)
    x += P;
  else if (x >= P)
    x -= P;
}

void setIO(string a) {
#ifndef LOCAL
  freopen((a + ".in").c_str(), "r", stdin);
  freopen((a + ".out").c_str(), "w", stdout);
#endif
}

template <typename T> inline void chkmin(T &a, const T &b) {
  if (a > b)
    a = b;
}
template <typename T> inline void chkmax(T &a, const T &b) {
  if (a < b)
    a = b;
}

LL X[N];
int cnt, n, col[N];
struct edge {
  int v;
  LL w;
  int pre;
} e[N * 2];
int dex, adj[N];

struct info {
  LL d1;
  int c1, go1;
  LL d2;
  int c2, go2;
  info operator+(const info &a) const {
    info ans = *this;
    if (a.d1 < ans.d1) {
      if (ans.c1 == a.c1) {
        ans.go1 = a.go1;
        ans.d1 = a.d1;
      } else {
        ans.d2 = ans.d1;
        ans.c2 = ans.c1;
        ans.go2 = ans.go1;
        ans.d1 = a.d1;
        ans.c1 = a.c1;
        ans.go1 = a.go1;
      }
    }
    if (a.d1 < ans.d2 && a.c1 != ans.c1) {
      ans.d2 = a.d1;
      ans.c2 = a.c1;
      ans.go2 = a.go1;
    }
    if (a.d2 < ans.d2 && a.c2 != ans.c1) {
      ans.d2 = a.d2;
      ans.c2 = a.c2;
      ans.go2 = a.go2;
    }
    return ans;
  }
} dp1[N], dp2[N], seq1[N], seq2[N], tmp[N];
int L[N], R[N], tot, go[N], ID[N], fa[N];

void dfs(int x, int p = 0) {
  int fir = 1;
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    dfs(v, x);
  }
  L[x] = tot + 1;
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    R[x] = ++tot;
    go[tot] = v;
    ID[v] = tot;
    tmp[tot] = dp1[v];
    tmp[tot].d1 += e[i].w;
    tmp[tot].d2 += e[i].w;
  }
  seq1[L[x]] = tmp[L[x]];
  seq2[R[x]] = tmp[R[x]];
  REP(i, L[x] + 1, R[x]) seq1[i] = seq1[i - 1] + tmp[i];
  PER(i, R[x] - 1, L[x]) seq2[i] = seq2[i + 1] + tmp[i];
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    if (fir) {
      info ww = dp1[v];
      ww.d1 += e[i].w;
      ww.d2 += e[i].w;
      dp1[x] = dp1[x] + ww;
      fir = 0;
    } else {
      info ww = dp1[v];
      ww.d1 += e[i].w;
      ww.d2 += e[i].w;
      dp1[x] = dp1[x] + ww;
    }
  }
}

int nxt[N];
LL ww[N];

LL dd[10005];
int dst = 0;
void _dfs(int x, int y) {
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == y)
      continue;
    dd[v] = dd[x] + e[i].w;
    _dfs(v, x);
  }
}
int chk(int u, int v, LL w) {
  REP(i, 1, n) dd[i] = 0;
  _dfs(u, 0);
  assert(dd[v] == w - X[v]);
}

void dfs2(int x, int p = 0) {
  dp1[x] = dp1[x] + dp2[x];
  info uu = dp1[x];
  //	assert(dp1[x].c1!=dp1[x].c2);
  //	assert(uu.d1<=uu.d2);
  //	chk(x,dp1[x].go1,dp1[x].d1);
  //	chk(x,dp1[x].go2,dp1[x].d2);
  if (uu.c1 == col[x]) {
    if (uu.d2 + X[x] < ww[col[x]]) {
      ww[col[x]] = uu.d2 + X[x];
      nxt[col[x]] = uu.c2;
    }
  } else {
    if (uu.d1 + X[x] < ww[col[x]]) {
      ww[col[x]] = uu.d1 + X[x];
      nxt[col[x]] = uu.c1;
    }
  }
  for (int i = adj[x]; i; i = e[i].pre) {
    int v = e[i].v;
    if (v == p)
      continue;
    info now = dp2[x];
    if (ID[v] != L[x])
      now = now + seq1[ID[v] - 1];
    if (ID[v] != R[x])
      now = now + seq2[ID[v] + 1];
    now.d1 += e[i].w;
    now.d2 += e[i].w;
    dp2[v] = dp2[v] + now;
    dfs2(v, x);
  }
}

int get(int x) {
  if (fa[x] == x)
    return x;
  return fa[x] = get(fa[x]);
}

int app[N];

int main() {
#ifdef LOCAL
  freopen("1.in", "r", stdin);
//	freopen("1.out", "w", stdout);
#endif

  scanf("%d", &n);
  REP(i, 1, n) scanf("%lld", X + i);
  REP(i, 1, n - 1) {
    int x, y, z;
    scanf("%d%d%d", &x, &y, &z);
    e[++dex] = (edge){y, z, adj[x]};
    adj[x] = dex;
    e[++dex] = (edge){x, z, adj[y]};
    adj[y] = dex;
  }

  REP(i, 1, n) col[i] = i;
  cnt = n;
  LL res = 0;
  while (cnt > 1) {
    REP(i, 1, cnt) ww[i] = 1e18;
    REP(i, 1, n) {
      dp1[i] = (info){X[i], col[i], i, (LL)1e18, 0, 0};
      dp2[i] = (info){X[i], col[i], i, (LL)1e18, 0, 0};
    }
    tot = 0;
    dfs(1, 0);
    dfs2(1, 0);
    REP(i, 1, cnt) fa[i] = i;
    REP(i, 1, cnt) {
      int u = get(i), v = get(nxt[i]);
      if (u == v)
        continue;
      res += ww[i];
      fa[u] = v;
    }
    REP(i, 1, n) col[i] = get(col[i]);
    REP(i, 1, cnt) app[i] = 0;
    cnt = 0;
    REP(i, 1, n) if (app[col[i]]) col[i] = app[col[i]];
    else col[i] = app[col[i]] = ++cnt;
  }

  cout << res << endl;

  return 0;
}