#include "robots.h"
#include <bits/stdc++.h>
using namespace std;
namespace r = ranges;
namespace v = ranges::views;
struct CompressedPQ {
priority_queue<int> pq;
vector<int> count;
CompressedPQ(int nmax = 0) : count(nmax) {}
void push(int i) {
if (count[i] == 0) pq.push(i);
++count.at(i);
}
int top() {
return pq.top();
}
int pop() {
int t = top();
--count[t];
if (count[t] == 0) pq.pop();
return t;
}
bool empty() {
return pq.empty();
}
};
template <typename R> vector<r::range_value_t<R>> vec_of(R &&range) {
return vector(range.begin(), range.end());
}
// return is in increasing order of y
vector<int> remove_one_dimension(const vector<pair<int, int>> &toys, int A, int B, int toys_per_bot) {
vector<vector<int>> per_xcoord(A + 1);
for (auto [x, y] : toys) per_xcoord[x].emplace_back(y);
CompressedPQ remaining(B + 1);
for (int i = 0; i <= A; ++i) {
/*cerr << "coord " << i << endl;*/
for (int y : per_xcoord[i]) {
/*cerr << "push " << y << endl;*/
remaining.push(y);
}
for (int to_remove = 0; i < A && to_remove < toys_per_bot && !remaining.empty(); ++to_remove) {
/*cerr << "pop robot " << i << endl;*/
remaining.pop();
}
}
vector<int> ans;
while (!remaining.empty()) {
ans.push_back(remaining.pop());
}
return ans;
}
bool can_put_away(const vector<pair<int, int>> &comp_toys, int A, int B, int toys_per_bot) {
vector<int> remaining_y = remove_one_dimension(comp_toys, A, B, toys_per_bot);
vector<int> toys_left(B + 1, 0);
for (int y : remaining_y) ++toys_left[y];
int total_toys_left = 0;
for (int i = 0; i <= B; ++i) {
total_toys_left += toys_left[i];
if (i < B) {
total_toys_left = max(0, total_toys_left - toys_per_bot);
}
}
return total_toys_left == 0;
}
int putaway(int A, int B, int T, int X[], int Y[], int W[], int S[]) {
vector<int> weight_limits(X, X + A);
vector<int> size_limits(Y, Y + B);
r::sort(weight_limits);
r::sort(size_limits);
// where is my ranges::to bwaaaaaaaaa
vector<pair<int, int>> toys = vec_of(
v::iota(0, T) | v::transform([&](int i) { return pair{W[i], S[i]}; }));
vector<pair<int, int>> compressed_toys = vec_of(
toys | v::transform([&](pair<int, int> p) {
return pair{(int)distance(weight_limits.begin(),
r::upper_bound(weight_limits, p.first)),
(int)distance(size_limits.begin(),
r::upper_bound(size_limits, p.second))};
}));
/*for (auto [x, y] : compressed_toys) cerr << "compressed " << x << " " << y << endl;*/
if (!can_put_away(compressed_toys, A, B, T))
return -1;
int lo = 1, hi = T;
while (lo < hi) {
int mid = (lo + hi) / 2;
/*cerr << "TRY " << mid << endl;*/
if (can_put_away(compressed_toys, A, B, mid))
hi = mid;
else
lo = mid + 1;
}
return lo;
}
