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

#include <bits/stdc++.h>
#define all(x) x.begin(), x.end()
using namespace std;
using ll = long long;
const int maxn = 1 << 20;
int n, k, q;
int x[maxn], l[maxn], r[maxn], w[maxn];
int res[maxn];
struct dsu {
  vector<int> p;
  vector<int> r, rtop;
  dsu(int n) : p(n), r(n, 1) {
    iota(all(p), 0);
    rtop = p;
  }
  int par(int v) { return p[v] == v ? v : p[v] = par(p[v]); }
  void unite(int u, int v) {
    u = par(u), v = par(v);
    if (u == v)
      return;
    int rr = rtop[u];
    if (r[u] < r[v])
      swap(u, v);
    p[v] = u;
    r[u] += r[v];
    rtop[u] = rr;
  }
};
ll d[maxn], s[maxn];
int par[maxn];
vector<int> g[maxn], p[maxn];
void dfs(int v) {
  d[v] += w[v];
  s[v] += 1;
  for (auto &i : g[v]) {
    d[i] = d[v];
    s[i] = s[v];
    par[i] = v;
    dfs(i);
  }
}
vector<ll> solve() {
  memset(d, 0, sizeof d);
  memset(s, 0, sizeof s);
  for (int i = 0; i < maxn; i++)
    g[i].clear(), p[i].clear();
  vector<ll> ans(q);
  memset(par, -1, sizeof par);
  for (int i = 0; i < n; i++) {
    auto it = lower_bound(x, x + n, x[i] + k) - x;
    if (it < n)
      g[it].push_back(i), par[i] = it;
  }
  for (int i = 0; i < n; i++)
    if (par[i] == -1)
      dfs(i);
  for (int i = 0; i < q; i++) {
    p[r[i]].push_back(i);
  }
  dsu f(n);
  for (int i = 0; i < n; i++) {
    for (auto j : g[i])
      f.unite(i, j);
    for (auto &qid : p[i]) {
      int x = l[qid];
      res[qid] = s[x];
      ans[qid] = d[x];
      int y = par[f.rtop[f.par(x)]];
      if (y != -1) {
        res[qid] -= s[y];
        ans[qid] -= d[y];
      }
    }
  }
  return ans;
}
int main() {
  cin.tie(0)->sync_with_stdio(0);
  cin >> n >> k;
  for (int i = 0; i < n; i++)
    cin >> x[i], w[i] = i;
  cin >> q;
  for (int i = 0; i < q; i++)
    cin >> l[i] >> r[i], l[i]--, r[i]--;
  auto a = solve();
  for (int i = 0; i < n; i++) {
    x[i] *= -1;
    l[i] = n - 1 - l[i];
    r[i] = n - 1 - r[i];
    swap(l[i], r[i]);
  }
  reverse(x, x + n);
  reverse(w, w + n);
  auto b = solve();
  for (int i = 0; i < q; i++)
    cout << b[i] - a[i] + res[i] << '\\n';
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define all(x) x.begin(), x.end()
using namespace std;
using ll = long long;
const int maxn = 1 << 20;
int n, k, q;
int x[maxn], l[maxn], r[maxn], w[maxn];
int res[maxn];
struct dsu {
  vector<int> p;
  vector<int> r, rtop;
  dsu(int n) : p(n), r(n, 1) {
    iota(all(p), 0);
    rtop = p;
  }
  int par(int v) { return p[v] == v ? v : p[v] = par(p[v]); }
  void unite(int u, int v) {
    u = par(u), v = par(v);
    if (u == v)
      return;
    int rr = rtop[u];
    if (r[u] < r[v])
      swap(u, v);
    p[v] = u;
    r[u] += r[v];
    rtop[u] = rr;
  }
};
ll d[maxn], s[maxn];
int par[maxn];
vector<int> g[maxn], p[maxn];
void dfs(int v) {
  d[v] += w[v];
  s[v] += 1;
  for (auto &i : g[v]) {
    d[i] = d[v];
    s[i] = s[v];
    par[i] = v;
    dfs(i);
  }
}
vector<ll> solve() {
  memset(d, 0, sizeof d);
  memset(s, 0, sizeof s);
  for (int i = 0; i < maxn; i++)
    g[i].clear(), p[i].clear();
  vector<ll> ans(q);
  for (int i = 0; i < n; i++) {
    auto it = lower_bound(x, x + n, x[i] + k) - x;
    if (it < n)
      g[it].push_back(i), par[i] = it;
  }
  memset(par, -1, sizeof par);
  for (int i = 0; i < n; i++)
    if (par[i] == -1)
      dfs(i);
  for (int i = 0; i < q; i++) {
    p[r[i]].push_back(i);
  }
  dsu f(n);
  for (int i = 0; i < n; i++) {
    for (auto j : g[i])
      f.unite(i, j);
    for (auto &qid : p[i]) {
      int x = l[qid];
      res[qid] = s[x];
      ans[qid] = d[x];
      int y = par[f.rtop[f.par(x)]];
      if (y != -1) {
        res[qid] -= s[y];
        ans[qid] -= d[y];
      }
    }
  }
  return ans;
}
int main() {
  cin.tie(0)->sync_with_stdio(0);
  cin >> n >> k;
  for (int i = 0; i < n; i++)
    cin >> x[i], w[i] = i;
  cin >> q;
  for (int i = 0; i < q; i++)
    cin >> l[i] >> r[i], l[i]--, r[i]--;
  auto a = solve();
  for (int i = 0; i < n; i++) {
    x[i] *= -1;
    l[i] = n - 1 - l[i];
    r[i] = n - 1 - r[i];
    swap(l[i], r[i]);
  }
  reverse(x, x + n);
  reverse(w, w + n);
  auto b = solve();
  for (int i = 0; i < q; i++)
    cout << b[i] - a[i] + res[i] << '\\n';
}