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

#include "bits/stdc++.h"
#include <cassert>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> pii;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
#define all(a) (a).begin(), (a).end()
#define vi vector<int>
#define pb push_back
#define INF (1LL << 59)
#define int ll

// verified AOJ GRL_1
#define MAX_V 100000
struct edge {
  int to, cost;
};
vector<edge> G[MAX_V];

void dijkstra(int s, vector<ll> &d) {
  priority_queue<pii, vector<pii>, greater<pii>> que;
  //    rep( i,d.size() )d[i]=INF;
  d[s] = 0;
  que.push(pii(0, s));

  while (que.size()) {
    pii p = que.top();
    que.pop();

    int v = p.second;
    if (d[v] < p.first)
      continue;

    rep(i, G[v].size()) {
      edge e = G[v][i];
      if (d[e.to] > d[v] + e.cost) {
        d[e.to] = d[v] + e.cost;
        que.push(pii(d[e.to], e.to));
      }
    }
  }
}

template <class T> class segtree {
private:
  vector<T> dat;
  int _size; // the number of leaves
  T _init;

public:
  segtree(int __size = 0, T init = T())
      : _size(__size), _init(init), dat(4 * __size, init) {
    if (__size == 0)
      return;
    int x = 1;
    while (x < _size)
      x *= 2;
    _size = x;
  }

  void init(int m, int x) {
    int i = m + _size;
    dat[i].pb(x);
  }

  void refresh() {
    for (int i = _size - 1; i > 0; i--) {
      T tmp = dat[2 * i];
      tmp.insert(tmp.end(), all(dat[2 * i + 1]));
      dat[i] = tmp;
    }
  }

  int find(int s, int t, int p, int num = 1, int a = 0, int b = -1) {
    if (b == -1)
      b = _size - 1; // I couldn't "int b=_size".
    if (b < s || t < a)
      return 0;
    if (s <= a && b <= t)
      return dat[num].size() -
             (lower_bound(all(dat[num]), p) - dat[num].begin());
    int ret, res;
    ret = find(s, t, p, num * 2, a, (a + b) / 2);
    res = find(s, t, p, num * 2 + 1, (a + b) / 2 + 1, b);

    return ret + res;
  }

  void sortAll() {
    rep(i, _size * 2 - 1) { sort(all(dat[i])); }
  }
};

signed main() {
  set<ll> valid_number;
  int n, m;
  cin >> n >> m;
  rep(i, m) {
    int a, b, c;
    scanf("%lld%lld%lld", &a, &b, &c);

    G[a].pb(edge{b, c});
    G[b].pb(edge{a, c});
  }

  vector<ll> ds(n, INF), dg(n, INF);
  dijkstra(0, ds);
  dijkstra(n - 1, dg);

  rep(i, ds.size()) valid_number.insert(ds[i]);

  int q;
  cin >> q;
  vector<pair<ll, ll>> query;
  rep(i, q) {
    ll a, b;
    scanf("%lld%lld", &a, &b);
    valid_number.insert(a);
    query.pb(pair<ll, ll>(a, b));
  }

  map<ll, int> conv;
  int c = 0;
  for (auto &num : valid_number) {
    conv[num] = c++;
  }

  segtree<vi> st(valid_number.size());
  rep(i, n) { st.init(conv[ds[i]], dg[i]); }
  st.refresh();
  st.sortAll();

  rep(i, q) {
    printf("%lld\\n", st.find(0, conv[query[i].first], query[i].second));
  }
}

Error: TLE

Fixed submission:

#include "bits/stdc++.h"
#include <cassert>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
#define all(a) (a).begin(), (a).end()
#define vi vector<int>
#define pb push_back
#define INF (1LL << 59)

// verified AOJ GRL_1
#define MAX_V 100000
struct edge {
  int to, cost;
};
vector<edge> G[MAX_V];

void dijkstra(int s, vector<ll> &d) {
  priority_queue<pii, vector<pii>, greater<pii>> que;
  //    rep( i,d.size() )d[i]=INF;
  d[s] = 0;
  que.push(pii(0, s));

  while (que.size()) {
    pii p = que.top();
    que.pop();

    int v = p.second;
    if (d[v] < p.first)
      continue;

    rep(i, G[v].size()) {
      edge e = G[v][i];
      if (d[e.to] > d[v] + e.cost) {
        d[e.to] = d[v] + e.cost;
        que.push(pii(d[e.to], e.to));
      }
    }
  }
}

// 0-index ????????????????°?????????(vector?????????, ?????????)
// (????´????°??????????¨sorted??????´??????????)
template <class T> class segtree {
private:
  vector<T> dat;
  int _size; // the number of leaves
  T _init;

public:
  segtree(int __size = 0, T init = T())
      : _size(__size), _init(init), dat(4 * __size, init) {
    if (__size == 0)
      return;
    int x = 1;
    while (x < _size)
      x *= 2;
    _size = x;
  }

  void update(int m, T x) {
    int i = m + _size;
    dat[i].insert(dat[i].end(), all(x));
    while (i != 1) {
      T tmp = dat[i];
      tmp.insert(tmp.end(), all(dat[i ^ 1]));
      dat[i / 2] = tmp;
      i /= 2;
    }
  }

  int find(int s, int t, int p, int num = 1, int a = 0, int b = -1) {
    if (b == -1)
      b = _size - 1; // I couldn't "int b=_size".
    if (b < s || t < a)
      return 0;
    if (s <= a && b <= t)
      return dat[num].size() -
             (lower_bound(all(dat[num]), p) - dat[num].begin());
    int ret, res;
    ret = find(s, t, p, num * 2, a, (a + b) / 2);
    res = find(s, t, p, num * 2 + 1, (a + b) / 2 + 1, b);

    return ret + res;
  }

  void sortAll() {
    rep(i, _size * 2 - 1) { sort(all(dat[i])); }
  }
};

int main() {
  set<ll> valid_number;
  int n, m;
  cin >> n >> m;
  rep(i, m) {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);

    G[a].pb(edge{b, c});
    G[b].pb(edge{a, c});
  }

  vector<ll> ds(n, INF), dg(n, INF);
  dijkstra(0, ds);
  dijkstra(n - 1, dg);

  rep(i, ds.size()) valid_number.insert(ds[i]);

  int q;
  cin >> q;
  vector<pii> query;
  rep(i, q) {
    int a, b;
    scanf("%d%d", &a, &b);
    valid_number.insert(a);
    query.pb(pii(a, b));
  }

  map<ll, int> conv;
  int c = 0;
  for (auto &num : valid_number) {
    conv[num] = c++;
  }

  segtree<vi> st(valid_number.size());
  rep(i, n) { st.update(conv[ds[i]], vi(1, dg[i])); }
  st.sortAll();

  rep(i, q) {
    printf("%d\\n", st.find(0, conv[query[i].first], query[i].second));
  }
}