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

#include <algorithm>
#include <stdio.h>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, pii> pi3;
typedef vector<pair<vector<int>, vector<pi3>>> inher;
vector<pi3> ans;
ll dat[200000];
#define INF 1000000000000000000LL
class tree {
public:
  // ここからいじらないところ(辺重みがあるときはうまくいじる)
  vector<int>
      idx; // 元々の index 木が小さくなるたびに実際の index は付け替えられる
  vector<vector<pii>> pat;
  vector<bool> flag;
  int num, cent; // num は頂点数、cent は重心の頂点番号 (0 <= cent < num)
  int dfs1(int node) {
    if (flag[node])
      return 0;
    flag[node] = true;
    int s = 1, maxi = 0;
    for (int i = 0; i < pat[node].size(); i++) {
      int r = dfs1(pat[node][i].first);
      s += r, maxi = max(maxi, r);
    }
    if (maxi <= num / 2 && (num - s) <= num / 2)
      cent = node;
    return s;
  }
  void init(vector<int> zid, vector<pi3> ed) {
    idx = zid;
    num = idx.size();
    flag.resize(num);
    pat.resize(num);
    for (int i = 0; i < ed.size(); i++) {
      pat[ed[i].second.first].push_back(
          make_pair(ed[i].second.second, ed[i].first));
      pat[ed[i].second.second].push_back(
          make_pair(ed[i].second.first, ed[i].first));
    }
    dfs1(0);
    fill(flag.begin(), flag.end(), false);
  }
  vector<int> sub1;
  vector<pi3> sub2;
  void dfs2(int node) {
    if (flag[node])
      return;
    flag[node] = true;
    int t = sub1.size();
    sub1.push_back(idx[node]);
    for (int i = 0; i < pat[node].size(); i++) {
      if (!flag[pat[node][i].first]) {
        sub2.push_back(
            make_pair(pat[node][i].second, make_pair(t, sub1.size())));
        dfs2(pat[node][i].first);
      }
    }
  }
  inher decompose() {
    inher r;
    fill(flag.begin(), flag.end(), false);
    flag[cent] = true;
    for (int i = 0; i < pat[cent].size(); i++) {
      sub1.clear();
      sub2.clear();
      dfs2(pat[cent][i].first);
      r.push_back(make_pair(sub1, sub2));
    }
    return r;
  }
  ll mn;
  int rr;
  void dfs3(int node, ll dep) {
    if (flag[node])
      return;
    flag[node] = true;
    if (mn > dep + dat[idx[node]])
      mn = dep + dat[idx[node]], rr = idx[node];
    ans.push_back(make_pair(dat[idx[cent]] + dat[idx[node]] + dep,
                            make_pair(idx[cent], idx[node])));
    for (int i = 0; i < pat[node].size(); i++)
      dfs3(pat[node][i].first, dep + pat[node][i].second);
  }
  void dfs4(int node, ll dep, int r) {
    if (flag[node])
      return;
    flag[node] = true;
    ans.push_back(make_pair(dat[idx[node]] + dep, make_pair(r, idx[node])));
    // printf("%d %d\\n",r, idx[node]);
    for (int i = 0; i < pat[node].size(); i++)
      dfs4(pat[node][i].first, dep + pat[node][i].second, r);
  }
  // ここまでいじらないところ
  void solve() // いじるところ
  {
    flag[cent] = true;
    ll m1 = INF, m2 = INF;
    int r1 = -1, r2 = -1;
    int t = -1;
    for (int i = 0; i < pat[cent].size(); i++) {
      mn = INF, rr = -1;
      dfs3(pat[cent][i].first, pat[cent][i].second);
      if (m1 > mn)
        m2 = m1, r2 = r1, m1 = mn, r1 = rr, t = i;
      else if (m2 > mn)
        m2 = mn, r2 = rr;
    }
    fill(flag.begin(), flag.end(), false);
    flag[cent] = true;
    // printf("   %lld %d\\n", m1, r1);
    for (int i = 0; i < pat[cent].size(); i++) {
      ll z;
      int r;
      if (t != i)
        z = m1, r = r1;
      else
        z = m2, r = r2;
      dfs4(pat[cent][i].first, z + pat[cent][i].second, r);
    }
  }
};
/*
重心分解
idx = {0,1,...,N-1}, ed は辺全部 として呼ぶ
tr.init で木が作られて重心が求まる
tr.solve で本質をやる
tr.decompose で分解を計算
子に対する solve の結果を利用するときは cetroiddecomposition の返り値として渡す

再帰するごとに頂点番号を 0 から頂点数 - 1 に取り換えるので、その取り換えた後の
index -> もとの index を表すのが idx ed は取り換え後の index で頂点を表現する
*/
void centroiddecomposition(vector<int> idx, vector<pi3> ed) {
  // for (int i = 0; i < idx.size(); i++)printf("%d ", idx[i] + 1);
  // printf("\\n");
  if (idx.size() <= 2) {
    if (idx.size() == 2)
      ans.push_back(make_pair(ed[0].first + dat[idx[0]] + dat[idx[1]],
                              make_pair(idx[0], idx[1])));
    return;
  }
  tree tr;
  tr.init(idx, ed);
  tr.solve();
  inher r = tr.decompose();
  tr.idx.clear();
  tr.pat.clear();
  tr.flag.clear();
  tr.sub1.clear();
  tr.sub2.clear();
  for (int i = 0; i < r.size(); i++)
    centroiddecomposition(r[i].first, r[i].second);
}
#define SIZE 200000
class unionfind {
public:
  int par[SIZE];
  int ran[SIZE];
  int ren[SIZE];
  void init() {
    for (int i = 0; i < SIZE; i++) {
      par[i] = i;
      ran[i] = 0;
      ren[i] = 1;
    }
  }
  int find(int a) {
    if (a == par[a])
      return a;
    else
      return par[a] = find(par[a]);
  }
  void unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a == b)
      return;
    if (ran[a] > ran[b]) {
      par[b] = a;
      ren[a] += ren[b];
    } else {
      par[a] = b;
      ren[b] += ren[a];
    }
    if (ran[a] == ran[b])
      ran[b]++;
  }
};
unionfind uf;
int main() {
  int num;
  scanf("%d", &num);
  vector<int> zzd;
  vector<pi3> ed;
  for (int i = 0; i < num; i++)
    scanf("%lld", &dat[i]), zzd.push_back(i);
  for (int i = 0; i < num - 1; i++) {
    int za, zb, zc;
    scanf("%d%d%d", &za, &zb, &zc);
    za--, zb--;
    ed.push_back(make_pair(zc, make_pair(za, zb)));
  }
  centroiddecomposition(zzd, ed);
  sort(ans.begin(), ans.end());
  uf.init();
  ll ret = 0;
  for (int i = 0; i < ans.size(); i++) {
    // printf("%lld %lld  %lld\\n", ans[i].second.first + 1, ans[i].second.second
    // + 1, ans[i].first);
    if (uf.find(ans[i].second.first) != uf.find(ans[i].second.second)) {
      ret += ans[i].first;
      uf.unite(ans[i].second.first, ans[i].second.second);
    }
  }
  printf("%lld\\n", ret);
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <stdio.h>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> pii;
typedef pair<ll, pii> pi3;
typedef vector<pair<vector<int>, vector<pi3>>> inher;
vector<pi3> ans;
ll dat[200000];
#define INF 1000000000000000000LL
class tree {
public:
  // ここからいじらないところ(辺重みがあるときはうまくいじる)
  vector<int>
      idx; // 元々の index 木が小さくなるたびに実際の index は付け替えられる
  vector<vector<pii>> pat;
  vector<bool> flag;
  int num, cent; // num は頂点数、cent は重心の頂点番号 (0 <= cent < num)
  int dfs1(int node) {
    if (flag[node])
      return 0;
    flag[node] = true;
    int s = 1, maxi = 0;
    for (int i = 0; i < pat[node].size(); i++) {
      int r = dfs1(pat[node][i].first);
      s += r, maxi = max(maxi, r);
    }
    if (maxi <= num / 2 && (num - s) <= num / 2)
      cent = node;
    return s;
  }
  void init(vector<int> zid, vector<pi3> ed) {
    idx = zid;
    num = idx.size();
    flag.resize(num);
    pat.resize(num);
    for (int i = 0; i < ed.size(); i++) {
      pat[ed[i].second.first].push_back(
          make_pair(ed[i].second.second, ed[i].first));
      pat[ed[i].second.second].push_back(
          make_pair(ed[i].second.first, ed[i].first));
    }
    dfs1(0);
    fill(flag.begin(), flag.end(), false);
  }
  vector<int> sub1;
  vector<pi3> sub2;
  void dfs2(int node) {
    if (flag[node])
      return;
    flag[node] = true;
    int t = sub1.size();
    sub1.push_back(idx[node]);
    for (int i = 0; i < pat[node].size(); i++) {
      if (!flag[pat[node][i].first]) {
        sub2.push_back(
            make_pair(pat[node][i].second, make_pair(t, sub1.size())));
        dfs2(pat[node][i].first);
      }
    }
  }
  inher decompose() {
    inher r;
    fill(flag.begin(), flag.end(), false);
    flag[cent] = true;
    for (int i = 0; i < pat[cent].size(); i++) {
      sub1.clear();
      sub2.clear();
      dfs2(pat[cent][i].first);
      r.push_back(make_pair(sub1, sub2));
    }
    return r;
  }
  ll mn;
  int rr;
  void dfs3(int node, ll dep) {
    if (flag[node])
      return;
    flag[node] = true;
    if (mn > dep + dat[idx[node]])
      mn = dep + dat[idx[node]], rr = idx[node];
    ans.push_back(make_pair(dat[idx[cent]] + dat[idx[node]] + dep,
                            make_pair(idx[cent], idx[node])));
    for (int i = 0; i < pat[node].size(); i++)
      dfs3(pat[node][i].first, dep + pat[node][i].second);
  }
  void dfs4(int node, ll dep, int r) {
    if (flag[node])
      return;
    flag[node] = true;
    ans.push_back(make_pair(dat[idx[node]] + dep, make_pair(r, idx[node])));
    // printf("%d %d\\n",r, idx[node]);
    for (int i = 0; i < pat[node].size(); i++)
      dfs4(pat[node][i].first, dep + pat[node][i].second, r);
  }
  // ここまでいじらないところ
  void solve() // いじるところ
  {
    flag[cent] = true;
    ll m1 = INF, m2 = INF;
    int r1 = -1, r2 = -1;
    int t = -1;
    for (int i = 0; i < pat[cent].size(); i++) {
      mn = INF, rr = -1;
      dfs3(pat[cent][i].first, pat[cent][i].second);
      if (m1 > mn)
        m2 = m1, r2 = r1, m1 = mn, r1 = rr, t = i;
      else if (m2 > mn)
        m2 = mn, r2 = rr;
    }
    fill(flag.begin(), flag.end(), false);
    flag[cent] = true;
    // printf("   %lld %d\\n", m1, r1);
    for (int i = 0; i < pat[cent].size(); i++) {
      ll z;
      int r;
      if (t != i)
        z = m1, r = r1;
      else
        z = m2, r = r2;
      dfs4(pat[cent][i].first, z + pat[cent][i].second, r);
    }
  }
};
/*
重心分解
idx = {0,1,...,N-1}, ed は辺全部 として呼ぶ
tr.init で木が作られて重心が求まる
tr.solve で本質をやる
tr.decompose で分解を計算
子に対する solve の結果を利用するときは cetroiddecomposition の返り値として渡す

再帰するごとに頂点番号を 0 から頂点数 - 1 に取り換えるので、その取り換えた後の
index -> もとの index を表すのが idx ed は取り換え後の index で頂点を表現する
*/
void centroiddecomposition(vector<int> idx, vector<pi3> ed) {
  // for (int i = 0; i < idx.size(); i++)printf("%d ", idx[i] + 1);
  // printf("\\n");
  if (idx.size() <= 2) {
    if (idx.size() == 2)
      ans.push_back(make_pair(ed[0].first + dat[idx[0]] + dat[idx[1]],
                              make_pair(idx[0], idx[1])));
    return;
  }
  tree tr;
  tr.init(idx, ed);
  tr.solve();
  inher r = tr.decompose();
  tr.idx.clear();
  tr.pat.clear();
  tr.flag.clear();
  tr.sub1.clear();
  tr.sub2.clear();
  for (int i = 0; i < r.size(); i++)
    centroiddecomposition(r[i].first, r[i].second);
}
#define SIZE 200000
class unionfind {
public:
  int par[SIZE];
  int ran[SIZE];
  int ren[SIZE];
  void init() {
    for (int i = 0; i < SIZE; i++) {
      par[i] = i;
      ran[i] = 0;
      ren[i] = 1;
    }
  }
  int find(int a) {
    if (a == par[a])
      return a;
    else
      return par[a] = find(par[a]);
  }
  void unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a == b)
      return;
    if (ran[a] > ran[b]) {
      par[b] = a;
      ren[a] += ren[b];
    } else {
      par[a] = b;
      ren[b] += ren[a];
    }
    if (ran[a] == ran[b])
      ran[b]++;
  }
};
unionfind uf;
int main() {
  int num;
  scanf("%d", &num);
  vector<int> zzd;
  vector<pi3> ed;
  for (int i = 0; i < num; i++)
    scanf("%lld", &dat[i]), zzd.push_back(i);
  for (int i = 0; i < num - 1; i++) {
    int za, zb, zc;
    scanf("%d%d%d", &za, &zb, &zc);
    za--, zb--;
    ed.push_back(make_pair(zc, make_pair(za, zb)));
  }
  centroiddecomposition(zzd, ed);
  sort(ans.begin(), ans.end());
  uf.init();
  ll ret = 0;
  for (int i = 0; i < ans.size(); i++) {
    // printf("%lld %lld  %lld\\n", ans[i].second.first + 1, ans[i].second.second
    // + 1, ans[i].first);
    if (uf.find(ans[i].second.first) != uf.find(ans[i].second.second)) {
      ret += ans[i].first;
      uf.unite(ans[i].second.first, ans[i].second.second);
    }
  }
  printf("%lld\\n", ret);
}