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

/*


*/
#include <algorithm>
#include <cctype>
#include <cstdio>
// #define gc() getchar()
#define MAXIN 400000
#define gc()                                                                   \\
  (SS == TT && (TT = (SS = IN) + fread(IN, 1, MAXIN, stdin), SS == TT)         \\
       ? EOF                                                                   \\
       : *SS++)
typedef long long LL;
const int N = 1e5 + 5, INF = 1 << 30;

int Ans, Enum, H[N], nxt[N << 1], to[N << 1], len[N << 1], A[N], Min, root,
    sz[N], cnt1, cnt2;
bool vis[N];
char IN[MAXIN], *SS = IN, *TT = IN;
struct Node {
  int s, dep, bel;
  bool operator<(const Node &x) const {
    return s < x.s; // 按s排序啊（我怎么按dep了mdzz）
  }
} f[N], g[N];

inline int read() {
  int now = 0;
  register char c = gc();
  for (; !isdigit(c); c = gc())
    ;
  for (; isdigit(c); now = now * 10 + c - 48, c = gc())
    ;
  return now;
}
inline void AE(int w, int u, int v) {
  to[++Enum] = v, nxt[Enum] = H[u], H[u] = Enum, len[Enum] = w;
  to[++Enum] = u, nxt[Enum] = H[v], H[v] = Enum, len[Enum] = w;
}
void FindRoot(int x, int fa, int tot) {
  int mx = 0;
  sz[x] = 1;
  for (int i = H[x], v; i; i = nxt[i])
    if (!vis[v = to[i]] && v != fa)
      FindRoot(v, x, tot), sz[x] += sz[v], mx = std::max(mx, sz[v]);
  mx = std::max(mx, tot - sz[x]);
  if (mx < Min)
    Min = mx, root = x;
}
void DFS(int x, int fa, int need, int need2, int sum, int dep,
         int anc) { // need:v->x当前点至少需要要有多少油量
                    // rest:v->x会剩下多少油量 need2:x->v最少需要多少起始油量
                    // sum:x->v一共多少油量 rest=A[root]+sum
  g[++cnt2] = (Node){need2, dep, anc}, sum += A[x];
  if (A[x] >= need)
    f[++cnt1] = (Node){sum, dep + 1, anc};
  for (int i = H[x], v, w; i; i = nxt[i])
    if (!vis[v = to[i]] && v != fa)
      w = len[i], DFS(v, x, w + std::max(0, need - A[x]),
                      std::max(need2, w - sum), sum - w, dep + 1, anc);
}
void Solve(int x) {
  vis[x] = 1, cnt1 = cnt2 = 1, f[1] = (Node){0, 1, 0}, g[1] = (Node){0, 0, 0};
  for (int i = H[x], v; i; i = nxt[i])
    if (!vis[v = to[i]])
      DFS(v, x, len[i], len[i], -len[i], 1, v);
  for (int i = 1, a = A[x]; i <= cnt1; ++i)
    f[i].s += a;
  std::sort(f + 1, f + 1 + cnt1), std::sort(g + 1, g + 1 + cnt2);
  Node a = (Node){0, 0, 0}, b = a;
  int res = 0;
  g[cnt2 + 1].s = INF;
  for (int i = 1, j = 1; i <= cnt1; ++i) {
    while (g[j].s <= f[i].s) {
      if (g[j].dep > a.dep)
        if (g[j].bel != a.bel)
          b = a, a = g[j];
        else
          a = g[j];
      else if (g[j].dep > b.dep && g[j].bel != a.bel)
        b = g[j];
      ++j;
    }
    if (f[i].bel != a.bel)
      res = std::max(res, f[i].dep + a.dep);
    else
      res = std::max(res, f[i].dep + b.dep);
  }
  Ans = std::max(Ans, res);
  for (int i = H[x], v; i; i = nxt[i])
    if (!vis[v = to[i]])
      Min = N, FindRoot(v, x, sz[v]), Solve(root);
}

int main() {
  int n = read();
  for (int i = 1; i <= n; ++i)
    A[i] = read();
  for (int i = 1; i < n; ++i)
    AE(read(), read(), read());
  Ans = 1, Min = N, FindRoot(1, 1, n), Solve(root), printf("%d\\n", Ans);

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

const int N = 1e6 + 8;
const int INF = 1e9 + 7;
struct edge {
  int u, v, w, nt;
} eg[N << 2];
int lt[N], sum;
int n;
int gas[N];
int vis[N];
int a[N], b[N];
int tot, root;
int size[N], sizeMax[N];
struct node {
  int x, len;
};
vector<node> f, g;
int p[N], num;
int ans, flag;
void add(int u, int v, int w) {
  eg[++sum] = (edge){u, v, w, lt[u]};
  lt[u] = sum;
}
struct Sement_Tree {
  int a[N], len;
  void init(int l) {
    len = l;
    for (int i = 0; i <= l; ++i)
      a[i] = -1;
  }
  void update(int x, int val) {
    x = len - x + 1;
    for (int i = x; i <= len; i += i & (-i))
      a[i] = max(a[i], val);
  }
  int query(int x) {
    x = len - x + 1;
    int res = -1;
    for (int i = x; i; i -= i & (-i))
      res = max(res, a[i]);
    return res;
  }
} T;
void init() {
  cin >> n;
  for (int i = 1; i <= n; ++i)
    cin >> gas[i];
  for (int i = 1; i < n; ++i) {
    int u, v, w;
    cin >> u >> v >> w;
    add(u, v, w);
    add(v, u, w);
  }
}
void getSize(int u, int fa = 0) {
  size[u] = 1;
  for (int i = lt[u]; i; i = eg[i].nt) {
    int v = eg[i].v;
    if (v == fa || vis[v])
      continue;
    getSize(v, u);
    size[u] += size[v];
  }
}
void getRoot(int u, int fa = 0) {
  for (int i = lt[u]; i; i = eg[i].nt) {
    int v = eg[i].v;
    if (v == fa || vis[v])
      continue;
    getRoot(v, u);
    sizeMax[u] = max(sizeMax[u], size[v]);
  }
  sizeMax[u] = max(sizeMax[u], tot - size[u]);
  if (sizeMax[u] < sizeMax[root])
    root = u;
}
void getToDp(int u, int fa = 0, int need = 0, int deep = 0) {
  b[u] = max(b[fa], need);
  g.push_back((node){b[u], deep});
  for (int i = lt[u]; i; i = eg[i].nt) {
    int v = eg[i].v;
    if (v == fa || vis[v])
      continue;
    getToDp(v, u, need + eg[i].w - gas[u], deep + 1);
  }
}
void getFromDp(int u, int fa = 0, int delta = 0, int deep = 0,
               int require = 0) {
  if (gas[u] >= require)
    a[u] = -delta;
  else
    a[u] = -1;
  if (a[u] >= 0)
    f.push_back((node){a[u], deep});
  for (int i = lt[u]; i; i = eg[i].nt) {
    int v = eg[i].v;
    if (v == fa || vis[v])
      continue;
    getFromDp(v, u, delta + eg[i].w - gas[v], deep + 1,
              max(0, require - gas[u]) + eg[i].w);
  }
}
vector<edge> vec;

void work(int u) {
  f.clear();
  getFromDp(u);
  num = 0;
  for (auto w : f)
    p[++num] = w.x;
  sort(p + 1, p + num + 1);
  num = unique(p + 1, p + num + 1) - p - 1;
  vec.clear();
  for (int i = lt[u]; i; i = eg[i].nt)
    if (!vis[eg[i].v])
      vec.push_back(eg[i]);
  b[u] = 0;
  for (int cas = 1; cas <= 2; ++cas) {
    T.init(num);
    T.update(1, 0);
    if (cas == 2)
      reverse(vec.begin(), vec.end());
    for (auto q : vec) {
      f.clear();
      g.clear();
      getToDp(q.v, u, q.w - gas[u], 1);
      getFromDp(q.v, u, q.w - gas[q.v], 1, q.w);
      for (auto w : g) {
        int x = lower_bound(p + 1, p + num + 1, w.x) - p;
        if (x != num + 1) {
          int mx = T.query(x); // 后?
          if (mx != -1)
            ans = max(ans, mx + w.len);
        }
      }
      for (auto w : f) {
        int x = lower_bound(p + 1, p + num + 1, w.x) - p;
        T.update(x, w.len);
      }
    }
  }
  int x = lower_bound(p + 1, p + num + 1, 0) - p;
  int mx = T.query(x);
  if (mx != -1)
    ans = max(ans, mx);
  vis[u] = 1;
  for (int i = lt[u]; i; i = eg[i].nt) {
    int v = eg[i].v;
    if (vis[v])
      continue;
    getSize(v);
    tot = size[v];
    root = 0;
    if (tot < ans + 1)
      continue;
    getRoot(v);
    work(v);
  }
}
void solve() {
  getSize(1);
  tot = n;
  sizeMax[root = 0] = n;
  getRoot(1);
  ans = 0;
  work(root);
  cout << ans + 1 << endl;
}
int main() {
  cin.sync_with_stdio(0);
  init();
  solve();
}