Escape route

#include "escape_route.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;

// Struct for priority queue operations on index set [1, n]
// push(i, v) overwrites value at position i if one already exists
// decKey is faster but requires that the new key is smaller than the old one
struct Prique {
    const ll INF = 4 * (ll)1e18;
    vector<pair<ll, int>> data;
    const int n;

    Prique(int siz) : n(siz), data(2*siz, {INF, -1}) { data[0] = {-INF, -1}; }
    void push(int i, ll v) {
        data[i+n] = {v, (v >= INF ? -1 : i)};
        for (i += n; i > 1; i >>= 1) data[i>>1] = min(data[i], data[i^1]);
    }
    void decKey(int i, ll v) {
        for (int j = i+n; data[j].first > v; j >>= 1) data[j] = {v, i};
    }
    pair<ll, int> pop() {
        auto res = data[1];
        push(res.second, INF);
        return res;
    }
};

const int N = 90;
const ll INF = 1e18;
ll tight_times[N][N][N];
ll wts[N][N], cs[N][N];
vector<pair<ll, int>> qs[N], ord[N];

vector<ll> calculate_necessary_time(
        int n, int m, ll s, int q, vector<int> A, vector<int> B,
        vector<ll> L, vector<ll> C, vector<int> U,
            vector<int> V,  vector<ll> T) {

    // N: number of cities, N <= 90
    // M: number of roads, M <= (N choose 2)
    // S: seconds in a day, S <= 1e15
    // Q: #queries, Q <= 3e6
    // A, B, L, C: road i connects A[i] and B[i], takes L[i] seconds to traverse, and closes at time C[i]
    // U, V, T: member j departs from city U[j] to city V[j] at time T[j]

    // Clear data
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            wts[i][j] = 0;
            cs[i][j] = -1;
        }
        cs[i][i] = 0;
    }
    for (int i = 0; i < n; ++i) {
        qs[i].clear();
        qs[i].shrink_to_fit();
        ord[i].clear();
        ord[i].shrink_to_fit();
    }

    // Assume complete graph
    for (int j = 0; j < m; ++j) {
        int a = A[j];
        int b = B[j];
        wts[a][b] = L[j];
        wts[b][a] = L[j];
        cs[a][b] = C[j] - L[j];
        cs[b][a] = C[j] - L[j];
    }

    // Calculate tight times : time to get from i to j, if we start at the last time we can take edge from i to x
    for (int i = 0; i < n; ++i) {
        for (int x = 0; x < n; ++x) {
            for (int j = 0; j < n; ++j) tight_times[i][x][j] = INF; // Clear tight times
            if (cs[i][x] == -1) continue;

            tight_times[i][x][i] = cs[i][x];

            Prique que(n);
            for (int j = i; j != -1; j = que.pop().second) {
                ll base = tight_times[i][x][j];
                ll sec = base % s;
                ll nxt = base + (s - sec);
                for (int t = 0; t < n; ++t) {
                    ll c = cs[j][t];
                    if (c == -1) continue;

                    ll off = (c < sec ? nxt : base) + wts[j][t];
                    if (off < tight_times[i][x][t]) {
                        tight_times[i][x][t] = off;
                        que.decKey(t, off);
                    }
                }
            }
        }
    }

    for (int j = 0; j < q; ++j) {
        qs[U[j]].emplace_back(T[j], j);
    }

    for (int i = 0; i < n; ++i) {
        sort(qs[i].begin(), qs[i].end());
        reverse(qs[i].begin(), qs[i].end());

        ord[i].resize(n);
        for (int j = 0; j < n; ++j) ord[i][j] = make_pair(cs[i][j], j);

        sort(ord[i].begin(), ord[i].end());
        reverse(ord[i].begin(), ord[i].end());
    }

    vector<ll> ans(q, INF);
    for (int i = 0; i < n; ++i) {
        vector<int> inds(n, 0);
        vector<ll> cur_dist(n, INF);
        cur_dist[i] = 0;

        int qi = 0;

        Prique que(n);
        que.decKey(i, -ord[i][inds[i]].first);
        while(true) {
            auto pr = que.pop();
            ll next_time = -pr.first;

            // Answer queries
            while(qi < qs[i].size() && qs[i][qi].first > next_time) {
                int ind = qs[i][qi].second;

                // Answer if possible in one day
                ans[ind] = min(ans[ind], cur_dist[V[ind]]);

                // Answer if need multiple days
                for (int t = 0; t < n; ++t) {
                    if (cur_dist[t] <= s) ans[ind] = min(ans[ind], s - T[ind] + tight_times[t][t][V[ind]]);
                }

                ++qi;
            }

            if (next_time < 0) break;

            // Update distances
            int j = pr.second;
            int x = ord[j][inds[j]].second;
            ++inds[j];

            if (inds[j] < n) que.decKey(j, -(ord[j][inds[j]].first - cur_dist[j]));

            for (int t = 0; t < n; ++t) {
                if (tight_times[j][x][t] >= s) continue;

                ll off = cur_dist[j] + tight_times[j][x][t] - cs[j][x];
                if (off < cur_dist[t]) {
                    cur_dist[t] = off;
                    if (inds[t] < n) {
                        que.decKey(t, -min(next_time, ord[t][inds[t]].first - off));
                    }
                }
            }
        }
    }
    return ans;
}