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

#include <bits/stdc++.h>

#define N 100010
#define inf 0x3f3f3f3f

using namespace std;

struct lhy {
  int x, y, l, next;
} edge[N << 1];

int tot, ans, n, l, nowsize, root, g[N], flag[N], nowl, newl;
pair<long long, int> q[N], Q[N];
int son[N], size[N], Max[N];

set<pair<long long, int>> A;
set<pair<long long, int>>::iterator it;

inline void adde(int x, int y, int l) {
  edge[++tot].x = x;
  edge[tot].y = y;
  edge[tot].l = l;
  edge[tot].next = son[x];
  son[x] = tot;
}

void GetRoot(int x, int fa) {
  size[x] = 1;
  Max[x] = 0;
  for (int i = son[x]; i; i = edge[i].next) {
    int y = edge[i].y;
    if (flag[y] || y == fa)
      continue;
    GetRoot(y, x);
    size[x] += size[y];
    Max[x] = max(Max[x], size[y]);
  }
  Max[x] = max(Max[x], nowsize - size[x]);
  if (Max[x] < Max[root])
    root = x;
}

void ask(long long v, int deep) {
  it = A.upper_bound(make_pair(v, inf));
  it--;
  if ((*it).first == -1ll * inf * inf)
    return;
  // printf("%d %d\\n", (*it).first, v);
  ans = max(ans, deep + (*it).second - 1);
}
int CNT;
void add(long long v, int deep) {
  it = A.upper_bound(make_pair(v, inf));
  it--;
  if ((*it).second >= deep)
    return;
  it++;
  set<pair<long long, int>>::iterator it2 = it, it1;
  for (; it2 != A.end(); it2++) {
    if ((*it2).second > deep)
      break;
  }
  while (it != it2) {
    it1 = it;
    it++;
    A.erase(it1);
    CNT++;
  }
  A.insert(make_pair(v, deep));
}

void dfs(int x, int fa, long long sumg, long long suml, int deep,
         long long nowmin1, long long nowmin2) {
  long long nowv = sumg - suml;
  // printf("%d %lld %lld %lld\\n", x, nowmin1, nowmin2, nowv);
  if (nowmin2 >= 0) {
    ans = max(ans, nowl + deep - 1);
    ask(nowv, deep);
    // printf("!%d\\n", ans);
  }
  if (nowmin1 < 0) {
    q[++l] = make_pair(-nowmin1, deep);
  } else {
    newl = max(newl, deep);
  }
  for (int i = son[x]; i; i = edge[i].next) {
    int y = edge[i].y;
    if (flag[y] || y == fa)
      continue;
    dfs(y, x, sumg + g[y], suml + edge[i].l, deep + 1,
        min(nowmin1, sumg - suml + g[root] - edge[i].l),
        min(nowmin2 + g[y] - edge[i].l, 0ll + g[y] - edge[i].l));
  }
}

void work(int x) {
  int L = 0;
  nowl = 1;
  for (int i = son[x]; i; i = edge[i].next) {
    int y = edge[i].y;
    newl = 1;
    l = 0;
    if (flag[y])
      continue;
    dfs(y, x, g[y], edge[i].l, 2, g[x] - edge[i].l, g[y] - edge[i].l);
    CNT = 0;
    // printf("%d\\n", l);
    for (int j = 1; j <= l; j++)
      add(q[j].first, q[j].second);
    // printf("%d\\n", CNT);
    // printf("%d\\n", A.size());
    Q[++L] = make_pair(y, edge[i].l);
    nowl = max(nowl, newl);
    ans = max(ans, nowl);
  }
  A.clear();
  A.insert(make_pair(-1ll * inf * inf, -inf));
  A.insert(make_pair(1ll * inf * inf, inf));
  nowl = 1;
  for (int i = L; i; i--) {
    int y = Q[i].first;
    newl = 1;
    l = 0;
    dfs(y, x, g[y], Q[i].second, 2, g[x] - Q[i].second, g[y] - Q[i].second);
    for (int j = 1; j <= l; j++)
      add(q[j].first, q[j].second);
    nowl = max(nowl, newl);
    ans = max(ans, nowl);
  }
  A.clear();
  A.insert(make_pair(-1ll * inf * inf, -inf));
  A.insert(make_pair(1ll * inf * inf, inf));
}

void Divide(int x) {
  flag[x] = 1;
  // printf("--%d\\n", x);
  work(x);
  // printf("%d\\n", ans);
  for (int i = son[x]; i; i = edge[i].next) {
    int y = edge[i].y;
    if (flag[y])
      continue;
    root = 0;
    nowsize = size[y];
    GetRoot(y, 0);
    Divide(root);
  }
  // printf("--%d\\n", x);
}

int main() {
  A.insert(make_pair(-1ll * inf * inf, -inf));
  A.insert(make_pair(1ll * inf * inf, inf));
  scanf("%d", &n);
  for (int i = 1; i <= n; i++)
    scanf("%d", &g[i]);
  for (int i = 1; i < n; i++) {
    int x, y, l;
    scanf("%d%d%d", &x, &y, &l);
    adde(x, y, l);
    adde(y, x, l);
  }
  ans = 1;
  Max[0] = n + 1;
  root = 0;
  nowsize = n;
  GetRoot(1, 0);
  Divide(root);
  printf("%d\\n", ans);
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define F(i, a, b) for (int i = (a); i <= (b); ++i)
using namespace std;

typedef long long ll;
typedef pair<int, int> P;
inline int RT(int l, int r) { return l + r | l != r; }
#define mst(a, b) memset(a, b, sizeof(a))
#define ___ freopen("c:\\\\code\\\\in.txt", "r", stdin);

const int N = 1e5 + 7;
int gas[N], n, mx, now, dep[N], ans, vis[N], in[N];
vector<P> g[N];
vector<int> head;
void dfs(int x, int fa) {
  dep[x] = dep[fa] + 1;
  for (auto &it : g[x])
    if (it.first != fa) {
      dfs(it.first, x);
    }
}

struct Node {
  int pre, gas, x, fa, val;
  Node(int a = 0, int b = 0, int c = 0, int d = 0, int e = 0)
      : pre(a), gas(b), x(c), fa(d), val(e) {}
} Q[N * 10];

P get(int x) {
  int l = 1, r = 1;
  Q[1] = Node(0, gas[x], x, 0, 1);
  Node now;
  int ansval = 0, idx = -1;
  while (l <= r) {
    now = Q[l++];
    if (now.val > ansval) {
      ansval = now.val, idx = l - 1;
      if (ansval == mx)
        return P(ansval, idx);
    }
    for (auto &it : g[now.x])
      if (it.first != now.fa && now.gas >= it.second) {
        Q[++r] = Node(l - 1, now.gas - it.second + gas[it.first], it.first,
                      now.x, now.val + 1);
      }
  }
  return P(ansval, idx);
}

int main() {
  scanf("%d", &n);
  F(i, 1, n) scanf("%d", gas + i);
  F(i, 2, n) {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);
    g[a].push_back(P(b, c));
    in[a]++, in[b]++;
    g[b].push_back(P(a, c));
  }
  dfs(1, 0);
  now = 0, mx = 0;
  F(i, 1, n) if (dep[i] > mx) mx = dep[i], now = i;
  dfs(now, 0);
  now = 0, mx = 0;
  F(i, 1, n) if (dep[i] > mx) mx = dep[i], now = i;
  int num = 200;
  srand(time(0));
  F(i, 1, num) head.push_back(rand() % n + 1);
  F(i, 1, n) if (in[i] == 2) head.push_back(i);
  for (auto &i : head) {
    P tmp = get(i);
    if (tmp.first > ans) {
      ans = tmp.first;
      // int x=tmp.second;
      // while(Q[x].pre)
      // {
      // x=Q[x].pre;
      // vis[Q[x].x]=1;
      // }
    }
    if (ans == mx)
      break;
  }
  printf("%d\\n", ans);
  return 0;
}