convex/nth_element_from_sorted_matrix.hpp

• View this file on GitHub
• Last update: 2025-07-04 07:32:29+09:00
• Include: #include "convex/nth_element_from_sorted_matrix.hpp"
• Link: https://uoj.ac/problem/891

Verified with

• test/1_mytest/nth_element_from_sorted_matrix.test.cpp

Code

// 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;
  }
  pq_min<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;
  }
  pq_min<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;
}

Back to top page