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

/*#include<cstdio>
#include<algorithm>
using namespace std;
int w[101000];
int n;
int main(){
    int i;

  //  freopen("/Users/joseunghyeon/Downloads/in (2)/eval_input.txt","r",stdin);
    int TC, TT;
    scanf("%d",&TC);
    for(TT=1;TT<=TC;TT++){
        scanf("%d",&n);
        for(i=1;i<=n;i++)scanf("%d",&w[i]);
        sort(w+1,w+n+1);
        printf("Case #%d\\n",TT);
        if(w[n]-w[1]>1){
            printf("No\\n");
            continue;
        }
        if(w[n]-w[1]==0){
            if(w[1] <= n/2 || w[1] == n-1)printf("Yes\\n");
            else printf("No\\n");
            continue;
        }
        for(i=1;i<n;i++)if(w[i]!=w[i+1])break;
        int t = i;
        if(t + 1 <= w[n] && t + (n-t)/2 >= w[n])printf("Yes\\n");
        else printf("No\\n");
    }
}
*/

/*#include<cstdio>
#include<algorithm>
using namespace std;
int n, w[201000];
int main(){
    char fp[200];

    int T[2]={6,18};
    freopen("/Users/joseunghyeon/Downloads/in (2)/eval_input.txt","w",stdout);
    printf("%d\\n",T[0]+T[1]);
    for(int TTT = 0; TTT < 2; TTT++){
        for(int TT=0;TT< T[TTT];TT++){
            fprintf(stderr,"%d %d\\n",TTT,TT);
            sprintf(fp, "/Users/joseunghyeon/Downloads/in (2)/%d_%02d.txt", TTT,
TT); freopen(fp, "r",stdin); int n, a; scanf("%d",&n); printf("%d\\n",n); for(int
i=0;i<n;i++){ scanf("%d",&a); printf("%d",a); if(i!=n-1)printf(" ");
            }
            printf("\\n");
        }
    }
    return 0;
}
*/

/*#include<cstdio>
#include<algorithm>
using namespace std;
int n, w[201000];
int main(){
    char fp[200];
    int TC = 0;
    int T[2]={6,18};
    freopen("/Users/joseunghyeon/Downloads/out (2)/eval_output.txt","w",stdout);
    for(int TTT = 0; TTT < 2; TTT++){
        for(int TT=0;TT< T[TTT];TT++){
            fprintf(stderr,"%d %d\\n",TTT,TT);
            sprintf(fp, "/Users/joseunghyeon/Downloads/out (2)/%d_%02d.txt",
TTT, TT); freopen(fp, "r",stdin); char rr[10]; scanf("%s",rr); printf("Case
#%d\\n%s\\n",++TC,rr);
        }
    }
    return 0;
}*/

/*
#include<algorithm>
#include<cstdio>
#include<map>
#include<vector>
using namespace std;
map<vector<int>,int>Map;
vector<int> Do(vector<int>T){
    vector<int>ret,T1,T2;
    if(T.size()==1){
        ret.push_back(T[0]+1);
        return ret;
    }
    int i;
    for(i=0;i<T.size()/2;i++){
        T1.push_back(min(T[i*2],T[i*2+1]));
        T2.push_back(max(T[i*2],T[i*2+1]));
    }
    vector<int>r1, r2;
    r1 = Do(T1);
    r2 = Do(T2);
    for(i=0;i<r1.size();i++)ret.push_back(r1[i]);
    for(i=0;i<r2.size();i++)ret.push_back(r2[i]);
    return ret;
}
int main(){
    int i, w[8];
    vector<int>TP(8);
    vector<int>res;
    for(i=0;i<8;i++)w[i]=i;
    do{
        for(i=0;i<8;i++)TP[i]=w[i];
        res = Do(TP);
        Map[res]  = 1;
    }while(next_permutation(w,w+8));
    for(auto &tp : Map){
        for(i=0;i<8;i++)printf("%d ",tp.first[i]);
        printf("\\n");
    }
}*/

#include <algorithm>
#include <cstdio>
#include <set>
#include <vector>
#define pli pair<long long, int>
#define N_ 201000
using namespace std;
vector<int> E[N_], L[N_];
int X[N_], n, C[N_], chk[N_], Num[N_][20];
long long D[N_], dis[N_][20], INF = 4e18;
int Root[N_][20];
vector<int> TP;
vector<pli> U[N_];
set<pli> Set;

int PV[N_], Mn[N_];
long long RR[N_];

bool MST[N_];

void DFS(int a, int pp) {
  int i;
  TP.push_back(a);
  C[a] = 1;
  for (i = 0; i < E[a].size(); i++) {
    if (chk[E[a][i]] || E[a][i] == pp)
      continue;
    D[E[a][i]] = D[a] + L[a][i];
    DFS(E[a][i], a);
    C[a] += C[E[a][i]];
  }
}
void Do(int a, int dep) {
  int i;
  D[a] = 0;
  TP.clear();
  DFS(a, 0);
  if (C[a] == 1)
    return;
  int M = C[a] + 1, mid = 0;
  for (i = 0; i < TP.size(); i++) {
    if (C[a] < C[TP[i]] * 2 && M > C[TP[i]]) {
      M = C[TP[i]];
      mid = TP[i];
    }
  }
  chk[mid] = 1;
  for (i = 0; i < E[mid].size(); i++) {
    if (!chk[E[mid][i]]) {
      Do(E[mid][i], dep + 1);
    }
  }
  chk[mid] = 0;
  TP.clear();
  D[mid] = 0;
  DFS(mid, 0);
  for (i = 0; i < TP.size(); i++) {
    Root[TP[i]][dep] = mid;
    dis[TP[i]][dep] = D[TP[i]];
    U[mid].push_back({D[TP[i]] + X[TP[i]], TP[i]});
  }
  sort(U[mid].begin(), U[mid].end());
  for (i = 0; i < U[mid].size(); i++) {
    Num[U[mid][i].second][dep] = i;
  }
  RR[mid] = INF;
  Mn[mid] = U[mid].size();
}

void Make(int a) {
  MST[a] = true;
  int i;
  for (i = 0; i < 20; i++) {
    if (!Root[a][i])
      break;
    int rt = Root[a][i], sz = U[rt].size();
    if (PV[rt] != sz && Mn[rt] != sz) {
      Set.erase(
          {U[rt][Mn[rt]].first + U[rt][PV[rt]].first, U[rt][PV[rt]].second});
    }
    Mn[rt] = min(Mn[rt], Num[a][i]);
    while (PV[rt] < sz && MST[U[rt][PV[rt]].second])
      PV[rt]++;
    if (PV[rt] != sz && Mn[rt] != sz) {
      Set.insert(
          {U[rt][Mn[rt]].first + U[rt][PV[rt]].first, U[rt][PV[rt]].second});
    }
  }
}

int main() {
  int i, a, b, c;
  scanf("%d", &n);
  for (i = 1; i <= n; i++) {
    scanf("%d", &X[i]);
  }
  for (i = 1; i < n; i++) {
    scanf("%d%d%d", &a, &b, &c);
    E[a].push_back(b);
    L[a].push_back(c);
    E[b].push_back(a);
    L[b].push_back(c);
  }
  Do(1, 0);

  Make(1);
  long long res = 0;
  for (int i = 0; i < n - 1; i++) {
    pli tp = *Set.begin();
    res += tp.first;
    Make(tp.second);
  }
  printf("%lld\\n", res);
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdio>
#include <set>
#include <vector>
#define pli pair<long long, int>
#define N_ 201000
using namespace std;
vector<int> E[N_], L[N_];
int X[N_], n, C[N_], chk[N_];
long long D[N_], dis[N_][20], INF = 4e18;
int Root[N_][20];
vector<int> TP;

multiset<long long> On[N_];
set<pli> Off[N_], Set;
bool MST[N_];

void DFS(int a, int pp) {
  int i;
  TP.push_back(a);
  C[a] = 1;
  for (i = 0; i < E[a].size(); i++) {
    if (chk[E[a][i]] || E[a][i] == pp)
      continue;
    D[E[a][i]] = D[a] + L[a][i];
    DFS(E[a][i], a);
    C[a] += C[E[a][i]];
  }
}
void Do(int a, int dep) {
  int i;
  D[a] = 0;
  TP.clear();
  DFS(a, 0);
  if (C[a] == 1)
    return;
  int M = C[a] + 1, mid = 0;
  for (i = 0; i < TP.size(); i++) {
    if (C[a] < C[TP[i]] * 2 && M > C[TP[i]]) {
      M = C[TP[i]];
      mid = TP[i];
    }
  }
  chk[mid] = 1;
  for (i = 0; i < E[mid].size(); i++) {
    if (!chk[E[mid][i]]) {
      Do(E[mid][i], dep + 1);
    }
  }
  chk[mid] = 0;
  TP.clear();
  D[mid] = 0;
  DFS(mid, 0);
  for (i = 0; i < TP.size(); i++) {
    Root[TP[i]][dep] = mid;
    dis[TP[i]][dep] = D[TP[i]];
    Off[mid].insert({D[TP[i]] + X[TP[i]], TP[i]});
  }

  Off[mid].insert({INF, -1});
  On[mid].insert(INF);
  Set.insert(
      {Off[mid].begin()->first + *On[mid].begin(), Off[mid].begin()->second});
}

void Make(int a) {
  MST[a] = true;
  int i;
  for (i = 0; i < 20; i++) {
    if (!Root[a][i])
      break;
    int rt = Root[a][i];
    pli r1 = {Off[rt].begin()->first + *On[rt].begin(),
              Off[rt].begin()->second};
    Off[rt].erase({dis[a][i] + X[a], a});
    On[rt].insert(dis[a][i] + X[a]);
    pli r2 = {Off[rt].begin()->first + *On[rt].begin(),
              Off[rt].begin()->second};
    Set.erase(r1);
    Set.insert(r2);
  }
}

int main() {
  int i, a, b, c;
  scanf("%d", &n);
  for (i = 1; i <= n; i++) {
    scanf("%d", &X[i]);
  }
  for (i = 1; i < n; i++) {
    scanf("%d%d%d", &a, &b, &c);
    E[a].push_back(b);
    L[a].push_back(c);
    E[b].push_back(a);
    L[b].push_back(c);
  }
  Do(1, 0);

  Make(1);
  long long res = 0;
  for (int i = 0; i < n - 1; i++) {
    pli tp = *Set.begin();
    res += tp.first;
    Make(tp.second);
  }
  printf("%lld\\n", res);
}