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

#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>

#define rep(i, n) for (int i = 0; i < (n); ++i)

using namespace std;

typedef long long ll;
typedef pair<ll, int> LP;
typedef vector<int> V;
typedef vector<ll> LV;

const ll inf = 2e18;

struct E {
  int to, cost;
};

int n;
int m;
vector<E> g[100000];
int q;
ll ds[100000];
ll dg[100000];
LV cp, cpg;
V st[1 << 18];

void add_edge(int a, int b, int c) {
  E e = {b, c};
  g[a].push_back(e);
  E f = {a, c};
  g[b].push_back(f);
}

void dijkstra(int s, ll *d) {
  fill_n(d, n, inf);

  priority_queue<LP, vector<LP>, greater<LP>> q;
  d[s] = 0;
  q.push(LP(0, s));
  while (!q.empty()) {
    LP v = q.top();
    q.pop();

    if (v.first > d[v.second]) {
      continue;
    }

    for (E &e : g[v.second]) {
      if (d[e.to] > v.first + e.cost) {
        d[e.to] = v.first + e.cost;
        q.push(LP(d[e.to], e.to));
      }
    }
  }
}

void compress(ll *ds, LV &cp) {
  rep(i, n) { cp.push_back(ds[i]); }
  sort(cp.begin(), cp.end());
  cp.erase(unique(cp.begin(), cp.end()), cp.end());

  rep(i, n) { ds[i] = lower_bound(cp.begin(), cp.end(), ds[i]) - cp.begin(); }
}

void merge(V a, V b, V &c) {
  int i = 0;
  int j = 0;
  while (i < a.size() || j < b.size()) {
    if (i < a.size()) {
      if (j < b.size() && a[i] > b[j]) {
        c.push_back(b[j++]);
      } else {
        c.push_back(a[i++]);
      }
    } else {
      c.push_back(b[j++]);
    }
  }
}

void init1() {
  rep(i, n) { st[(1 << 17) + ds[i] - 1].push_back(dg[i]); }
  rep(i, n) {
    sort(st[(1 << 17) + i - 1].begin(), st[(1 << 17) + i - 1].end());
  }
}

V init2(int l = 0, int r = 1 << 17, int k = 0) {
  if (r - l == 1) {
    return st[k];
  }
  merge(init2(l, (l + r) / 2, 2 * k + 1), init2((l + r) / 2, r, 2 * k + 2),
        st[k]);
  return st[k];
}

int count_dj(int fs, int fg, int l = 0, int r = 1 << 17, int k = 0) {
  if (fs <= l) {
    return 0;
  }

  if (r <= fs) {
    return st[k].end() - lower_bound(st[k].begin(), st[k].end(), fg);
  }

  int m = (l + r) / 2;
  int p = count_dj(fs, fg, l, m, 2 * k + 1);
  int q = count_dj(fs, fg, m, r, 2 * k + 2);
  return p + q;
}

int main() {
  cin >> n >> m;
  rep(i, m) {
    int a, b, c;
    cin >> a >> b >> c;
    add_edge(a, b, c);
  }

  dijkstra(0, ds);
  dijkstra(n - 1, dg);

  compress(ds, cp);
  compress(dg, cpg);

  init1();
  init2();

  cin >> q;
  rep(i, q) {
    ll fs, fg;
    cin >> fs >> fg;
    fs = upper_bound(cp.begin(), cp.end(), fs) - cp.begin();
    fg = lower_bound(cpg.begin(), cpg.end(), fg) - cpg.begin();
    cout << count_dj(fs, fg) << endl;
  }
  return 0;
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>

#define rep(i, n) for (int i = 0; i < (n); ++i)

using namespace std;

typedef long long ll;
typedef pair<ll, int> LP;
typedef vector<int> V;
typedef vector<ll> LV;

const ll inf = 2e18;

struct E {
  int to, cost;
};

int n;
int m;
vector<E> g[100000];
int q;
ll ds[100000];
ll dg[100000];
LV cp, cpg;
V dbl[17][100000];

void add_edge(int a, int b, int c) {
  E e = {b, c};
  g[a].push_back(e);
  E f = {a, c};
  g[b].push_back(f);
}

void dijkstra(int s, ll *d) {
  fill_n(d, n, inf);

  priority_queue<LP, vector<LP>, greater<LP>> q;
  d[s] = 0;
  q.push(LP(0, s));
  while (!q.empty()) {
    LP v = q.top();
    q.pop();

    if (v.first > d[v.second]) {
      continue;
    }

    for (E &e : g[v.second]) {
      if (d[e.to] > v.first + e.cost) {
        d[e.to] = v.first + e.cost;
        q.push(LP(d[e.to], e.to));
      }
    }
  }
}

void compress(ll *ds, LV &cp) {
  rep(i, n) { cp.push_back(ds[i]); }
  sort(cp.begin(), cp.end());
  cp.erase(unique(cp.begin(), cp.end()), cp.end());

  rep(i, n) { ds[i] = lower_bound(cp.begin(), cp.end(), ds[i]) - cp.begin(); }
}

void merge(V &a, V &b, V &c) {
  int i = 0;
  int j = 0;
  while (i < a.size() || j < b.size()) {
    if (i < a.size()) {
      if (j < b.size() && a[i] > b[j]) {
        c.push_back(b[j++]);
      } else {
        c.push_back(a[i++]);
      }
    } else {
      c.push_back(b[j++]);
    }
  }
}

void make_dbl() {
  rep(i, n) { dbl[0][ds[i]].push_back(dg[i]); }
  rep(i, n) { sort(dbl[0][i].begin(), dbl[0][i].end()); }

  for (int i = 1; i < 17; ++i) {
    rep(j, n) {
      int v = j + (1 << i - 1);
      if (v < n) {
        merge(dbl[i - 1][j], dbl[i - 1][v], dbl[i][j]);
      } else {
        dbl[i][j] = dbl[i - 1][j];
      }
    }
  }
}

int count_dj(int fs, int fg) {
  int s = upper_bound(cp.begin(), cp.end(), fs) - cp.begin();

  fg = lower_bound(cpg.begin(), cpg.end(), fg) - cpg.begin();

  int ans = 0;
  for (int i = 16; i >= 0; --i) {
    if (s >> i & 1) {
      s -= 1 << i;
      ans +=
          dbl[i][s].end() - lower_bound(dbl[i][s].begin(), dbl[i][s].end(), fg);
    }
  }
  return ans;
}

int main() {
  cin >> n >> m;
  rep(i, m) {
    int a, b, c;
    cin >> a >> b >> c;
    add_edge(a, b, c);
  }

  dijkstra(0, ds);
  dijkstra(n - 1, dg);

  compress(ds, cp);
  compress(dg, cpg);

  make_dbl();

  cin >> q;
  rep(i, q) {
    ll fs, fg;
    cin >> fs >> fg;
    cout << count_dj(fs, fg) << endl;
  }
  return 0;
}