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

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll INF = 1e17;
using P = pair<ll, ll>;

template <typename T> class SegmentTree {
  using func_t = function<T(const T &, const T &)>;
  const int n;
  const T id;
  func_t merge;
  vector<T> data;
  int size(int n) {
    int res;
    for (res = 1; res < n; res <<= 1)
      ;
    return res;
  }
  ll sub(int l, int r, int node, int lb, int ub, ll val) {
    if (ub <= l || r <= lb)
      return 0;
    if (l <= lb && ub <= r)
      return data[node].end() -
             lower_bound(data[node].begin(), data[node].end(), val);
    return sub(l, r, node * 2, lb, (lb + ub) / 2, val) +
           sub(l, r, node * 2 + 1, (lb + ub) / 2, ub, val);
  }

public:
  SegmentTree(const vector<T> &data_, T id_, func_t merge_)
      : n(size(data_.size())), id(id_), merge(merge_), data(n * 2, id_) {
    for (int i = 0; i < (int)data_.size(); i++)
      data[i + n] = data_[i];
    for (int i = n - 1; i >= 0; i--)
      data[i] = merge(data[i * 2], data[i * 2 + 1]);
  }
  ll Find(int l, int r, ll val) { return sub(l, r + 1, 1, 0, n, val); }
};

int main() {
  cin.sync_with_stdio(false);
  ll n, m;
  cin >> n >> m;
  vector<vector<P>> G(n);
  for (int i = 0, a, b, c; i < m; i++) {
    cin >> a >> b >> c;
    G[a].push_back(P(b, c));
    G[b].push_back(P(a, c));
  }
  vector<ll> sd(n, INF), td(n, INF);
  priority_queue<P, vector<P>, greater<P>> pq1;
  pq1.push(P(0, 0));
  while (!pq1.empty()) {
    auto p = pq1.top();
    pq1.pop();
    if (sd[p.second] != INF)
      continue;
    sd[p.second] = p.first;
    for (auto to : G[p.second]) {
      if (sd[to.first] == INF) {
        pq1.push(P(p.first + to.second, to.first));
      }
    }
  }
  priority_queue<P, vector<P>, greater<P>> pq2;
  pq2.push(P(0, n - 1));
  while (!pq2.empty()) {
    auto p = pq2.top();
    pq2.pop();
    if (td[p.second] != INF)
      continue;
    td[p.second] = p.first;
    for (auto to : G[p.second]) {
      if (td[to.first] == INF) {
        pq2.push(P(p.first + to.second, to.first));
      }
    }
  }
  vector<P> dd(n);
  for (int i = 0; i < n; i++) {
    dd[i].first = sd[i];
    dd[i].second = td[i];
  }
  sort(dd.begin(), dd.end());
  vector<vector<ll>> ba(n);
  for (int i = 0; i < n; i++) {
    sd[i] = dd[i].first;
    ba[i].push_back(dd[i].second);
  }
  SegmentTree<vector<ll>> st(ba, vector<ll>(),
                             [](const vector<ll> &l, const vector<ll> &r) {
                               int nl = l.size(), nr = r.size();
                               vector<ll> res;
                               res.reserve(nl + nr);
                               for (int i = 0, j = 0; i < nl || j < nr;) {
                                 if (j == nr) {
                                   res.push_back(l[i++]);
                                   continue;
                                 }
                                 if (i == nl) {
                                   res.push_back(r[j++]);
                                   continue;
                                 }
                                 if (l[i] <= r[j]) {
                                   res.push_back(l[i++]);
                                 } else {
                                   res.push_back(r[j++]);
                                 }
                               }
                               return res;
                             });
  ll q, fs, fg;
  cin >> q;
  while (q--) {
    cin >> fs >> fg;
    int l = 0, r = upper_bound(sd.begin(), sd.end(), fs) - sd.begin() - 1;
    printf("%lld\\n", st.Find(l, r, fg));
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll INF = 1e17;
using P = pair<ll, ll>;

template <typename T> class SegmentTree {
  using func_t = function<T(const T &, const T &)>;
  const int n;
  const T id;
  func_t merge;
  vector<T> data;
  int size(int n) {
    int res;
    for (res = 1; res < n; res <<= 1)
      ;
    return res;
  }
  T sub(int l, int r, int node, int lb, int ub) {
    if (ub <= l || r <= lb)
      return id;
    if (l <= lb && ub <= r)
      return data[node];
    return merge(sub(l, r, node * 2, lb, (lb + ub) / 2),
                 sub(l, r, node * 2 + 1, (lb + ub) / 2, ub));
  }

public:
  SegmentTree(const vector<T> &data_, T id_, func_t merge_)
      : n(size(data_.size())), id(id_), merge(merge_), data(n * 2, id_) {
    for (int i = 0; i < (int)data_.size(); i++)
      data[i + n] = data_[i];
    for (int i = n - 1; i >= 0; i--)
      data[i] = merge(data[i * 2], data[i * 2 + 1]);
  }
  T Find(int i) { return data[i + n]; }
  T Find(int l, int r) { return sub(l, r + 1, 1, 0, n); }
};

int main() {
  ll n, m;
  cin >> n >> m;
  vector<vector<P>> G(n);
  for (int i = 0, a, b, c; i < m; i++) {
    cin >> a >> b >> c;
    G[a].push_back(P(b, c));
    G[b].push_back(P(a, c));
  }
  vector<ll> sd(n, INF), td(n, INF);
  priority_queue<P, vector<P>, greater<P>> pq1;
  pq1.push(P(0, 0));
  while (!pq1.empty()) {
    auto p = pq1.top();
    pq1.pop();
    if (sd[p.second] != INF)
      continue;
    sd[p.second] = p.first;
    for (auto to : G[p.second]) {
      if (sd[to.first] == INF) {
        pq1.push(P(p.first + to.second, to.first));
      }
    }
  }
  priority_queue<P, vector<P>, greater<P>> pq2;
  pq2.push(P(0, n - 1));
  while (!pq2.empty()) {
    auto p = pq2.top();
    pq2.pop();
    if (td[p.second] != INF)
      continue;
    td[p.second] = p.first;
    for (auto to : G[p.second]) {
      if (td[to.first] == INF) {
        pq2.push(P(p.first + to.second, to.first));
      }
    }
  }
  vector<P> dd(n);
  for (int i = 0; i < n; i++) {
    dd[i].first = sd[i];
    dd[i].second = td[i];
  }
  sort(dd.begin(), dd.end());
  vector<vector<ll>> ba(n);
  for (int i = 0; i < n; i++) {
    sd[i] = dd[i].first;
    ba[i].push_back(dd[i].second);
  }
  SegmentTree<vector<ll>> st(ba, vector<ll>(),
                             [](const vector<ll> &l, const vector<ll> &r) {
                               int nl = l.size(), nr = r.size();
                               vector<ll> res;
                               res.reserve(nl + nr);
                               for (int i = 0, j = 0; i < nl || j < nr;) {
                                 if (j == nr) {
                                   res.push_back(l[i++]);
                                   continue;
                                 }
                                 if (i == nl) {
                                   res.push_back(r[j++]);
                                   continue;
                                 }
                                 if (l[i] <= r[j]) {
                                   res.push_back(l[i++]);
                                 } else {
                                   res.push_back(r[j++]);
                                 }
                               }
                               return res;
                             });
  ll q, fs, fg;
  cin >> q;
  while (q--) {
    cin >> fs >> fg;
    int l = 0, r = upper_bound(sd.begin(), sd.end(), fs) - sd.begin() - 1;
    auto vv = st.Find(l, r);
    cout << (r - (lower_bound(vv.begin(), vv.end(), fg) - vv.begin()) + 1)
         << endl;
  }
  return 0;
}