#include "robots.h"
#include <bits/stdc++.h>
#include <pstl/glue_algorithm_defs.h>
using namespace std;
namespace r = ranges;
namespace v = ranges::views;
vector<pair<int, int>> remove_one_dim(vector<pair<int, int>> toys, vector<int> limits, int toys_per_bot) {
int ctoy = 0;
r::sort(toys);
r::sort(limits);
// toys are sorted by increasing weight
multiset<pair<int, int>> remaining_toys; // toys here are stored in (y, x) order
for (int weight_limit : limits) {
while (ctoy < toys.size() && toys[ctoy].first < weight_limit) {
auto toy = toys[ctoy++];
remaining_toys.emplace(toy.second, toy.first);
}
for (int _i : v::iota(0, toys_per_bot)) {
if (remaining_toys.empty()) break;
auto [y, x] = *prev(remaining_toys.end());
cerr << "limit " << weight_limit << " removes toy " << x << " " << y << endl;
remaining_toys.erase(prev(remaining_toys.end()));
}
}
while (ctoy < toys.size()) {
auto toy = toys[ctoy++];
remaining_toys.emplace(toy.second, toy.first);
}
toys = vector(remaining_toys.begin(), remaining_toys.end());
for (auto &t : toys) t = {t.second, t.first};
return toys;
}
bool can_put_away(vector<pair<int, int>> toys,
vector<int> weight_limits,
vector<int> size_limits,
int toys_per_bot) {
toys = remove_one_dim(toys, weight_limits, toys_per_bot);
for (auto &t : toys) t = {t.second, t.first};
cerr << "==" << endl;
toys = remove_one_dim(toys, size_limits, toys_per_bot);
cerr << "toys remaining: " << toys.size() << endl;
return toys.empty();
}
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);
// where is my ranges::to bwaaaaaaaaa
auto _toys =
v::iota(0, T) | v::transform([&](int i) { return pair{W[i], S[i]}; });
vector<pair<int, int>> toys(_toys.begin(), _toys.end());
if (!can_put_away(toys, weight_limits, size_limits, T)) return -1;
int lo = 1, hi = T;
while (lo < hi) {
int mid = (lo + hi) / 2;
cerr << "try " << mid << endl;
if (can_put_away(toys, weight_limits, size_limits, mid)) hi = mid;
else lo = mid + 1;
}
return lo;
}
