#include "overtaking.h"
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
struct Bus {
ll arrivalTime;
ll pace;
int id;
};
bool byArrivalThenPace(const Bus &a, const Bus &b) {
if (a.arrivalTime == b.arrivalTime) return a.pace < b.pace;
return a.arrivalTime < b.arrivalTime;
}
vector<Bus> buses;
vector<Bus> buses_init; // initial state for reset
vector<int> stations;
int N, M;
vector<ll> arrivalTimes_station[1005]; // per station sorted arrival times of buses
ll computeExpectedArrival(const Bus& b, int j) {
return b.arrivalTime + (stations[j] - stations[j - 1]) * b.pace;
}
void init(int Lp, int Np, vector<ll> Tp, vector<int> Wp, int Xp, int Mp, vector<int> Sp) {
buses.clear();
stations = Sp;
// Keep only buses with pace >= reserve pace Xp (filtering)
for (int i = 0; i < Np; i++) {
if (Wp[i] >= Xp)
buses.push_back({Tp[i], Wp[i], i});
}
// Add reserve bus with arrivalTime=0 (will set real departure time per query)
buses.push_back({0, Xp, (int)Np});
N = (int)buses.size();
M = Mp;
// Save initial state for reset before queries
buses_init = buses;
// Simulate station by station
for (int j = 1; j < M; j++) {
sort(buses.begin(), buses.end(), byArrivalThenPace);
ll curMax = 0;
vector<ll> arrivalTimes;
arrivalTimes.reserve(N);
for (int i = 0; i < N; i++) {
ll expected = computeExpectedArrival(buses[i], j);
buses[i].arrivalTime = max(expected, curMax);
curMax = max(curMax, buses[i].arrivalTime);
arrivalTimes.push_back(buses[i].arrivalTime);
}
arrivalTimes_station[j] = arrivalTimes;
}
}
ll arrival_time(ll Y) {
// Reset buses to initial state
buses = buses_init;
ll arrival = Y;
ll pace = 0;
// Set reserve bus departure time & get pace
for (auto &b : buses) {
if (b.id == N - 1) {
b.arrivalTime = Y;
pace = b.pace;
break;
}
}
// Simulate reserve bus arrival through stations with binary search
for (int j = 1; j < M; j++) {
const vector<ll>& arrTimes = arrivalTimes_station[j];
int pos = int(upper_bound(arrTimes.begin(), arrTimes.end(), arrival) - arrTimes.begin());
ll blockingArrival = (pos == 0) ? 0 : arrTimes[pos - 1];
ll dist = stations[j] - stations[j - 1];
arrival = max(arrival + dist * pace, blockingArrival);
// Update reserve bus arrival time in buses vector
for (auto &b : buses) {
if (b.id == N - 1) {
b.arrivalTime = arrival;
break;
}
}
}
// Return final arrival time of reserve bus
for (const auto &b : buses) {
if (b.id == N - 1)
return b.arrivalTime;
}
return -1; // Should never happen
}
