This documentation is automatically generated by online-judge-tools/verification-helper
#include "convex/nth_element_from_sorted_matrix.hpp"
// https://uoj.ac/problem/891
template <typename T, typename F>
vi nth_element_from_sorted_matrix(ll N, ll M, ll K, F f, int k1 = 0, int k2 = 0,
bool tr = false) {
if (K == 0) return vi(N, 0);
if (N > M) {
vi A = nth_element_from_sorted_matrix<T>(M, N, K, f, k2, k1, !tr);
vi B(N + 1);
FOR(i, M) B[0] += 1, B[A[i]] -= 1;
FOR(i, N) B[i + 1] += B[i];
B.pop_back();
return B;
}
vi A(N);
if (K > N) {
A = nth_element_from_sorted_matrix<T>(N, M / 2, (K - N) / 2, f, k1, k2 + 1,
tr);
for (auto &a: A) a *= 2;
K = K - (K - N) / 2 * 2;
}
pqg<pair<T, int>> que;
auto g = [&](ll i, ll j) -> T {
i = ((i + 1) << k1) - 1;
j = ((j + 1) << k2) - 1;
return (tr ? f(j, i) : f(i, j));
};
if (A[0] < M) que.emplace(g(0, A[0]), 0);
FOR(i, 1, N) if (A[i] < A[i - 1]) que.emplace(g(i, A[i]), i);
while (K) {
--K;
auto [x, i] = POP(que);
A[i]++;
if (K == 0) break;
if (A[i] < M && (i == 0 || A[i - 1] > A[i])) { que.emplace(g(i, A[i]), i); }
if (i + 1 < N && A[i + 1] == A[i] - 1) {
que.emplace(g(i + 1, A[i + 1]), i + 1);
}
}
return A;
}
#line 1 "convex/nth_element_from_sorted_matrix.hpp"
// https://uoj.ac/problem/891
template <typename T, typename F>
vi nth_element_from_sorted_matrix(ll N, ll M, ll K, F f, int k1 = 0, int k2 = 0,
bool tr = false) {
if (K == 0) return vi(N, 0);
if (N > M) {
vi A = nth_element_from_sorted_matrix<T>(M, N, K, f, k2, k1, !tr);
vi B(N + 1);
FOR(i, M) B[0] += 1, B[A[i]] -= 1;
FOR(i, N) B[i + 1] += B[i];
B.pop_back();
return B;
}
vi A(N);
if (K > N) {
A = nth_element_from_sorted_matrix<T>(N, M / 2, (K - N) / 2, f, k1, k2 + 1,
tr);
for (auto &a: A) a *= 2;
K = K - (K - N) / 2 * 2;
}
pqg<pair<T, int>> que;
auto g = [&](ll i, ll j) -> T {
i = ((i + 1) << k1) - 1;
j = ((j + 1) << k2) - 1;
return (tr ? f(j, i) : f(i, j));
};
if (A[0] < M) que.emplace(g(0, A[0]), 0);
FOR(i, 1, N) if (A[i] < A[i - 1]) que.emplace(g(i, A[i]), i);
while (K) {
--K;
auto [x, i] = POP(que);
A[i]++;
if (K == 0) break;
if (A[i] < M && (i == 0 || A[i - 1] > A[i])) { que.emplace(g(i, A[i]), i); }
if (i + 1 < N && A[i + 1] == A[i] - 1) {
que.emplace(g(i + 1, A[i + 1]), i + 1);
}
}
return A;
}