convex/nth_element_from_sorted_lists.hpp

• View this file on GitHub
• Last update: 2025-07-04 07:32:29+09:00
• Include: #include "convex/nth_element_from_sorted_lists.hpp"
• Link: https://codeforces.com/contest/1275/problem/F
• Link: https://codeforces.com/problemset/problem/1034/D

Verified with

• test/1_mytest/nth_element_from_sorted_lists.test.cpp

Code

// 行の要素数が S[0], S[1], ..., S[N-1]
// O(N(logN+logK)) time
// query num: O(NlogK), あまり最適じゃないと思う
// 通ったがあやしいっぽい: https://codeforces.com/contest/1275/problem/F
// 列方向もソートされているなら matrix を使おう!!
// https://codeforces.com/problemset/problem/1034/D
template <typename T, typename F>
vi nth_element_from_sorted_lists(vi S, ll K, F f, int k = 0) {
  ll N = len(S);
  ll sm = 0;
  for (auto& x : S) sm += x >> k;
  assert(0 <= K && K <= sm);
  if (K == 0) return vi(N, 0);
  if (K == sm) return S;

  ll row = 0;
  for (auto& x : S) row += (x >= (1LL << k));

  auto g = [&](int i, ll j) -> T {
    j = ((j + 1) << k) - 1;
    return (j >= S[i] ? infty<T> : f(i, j));
  };
  vi A(N);
  if (K > row) {
    A = nth_element_from_sorted_lists<T>(S, (K - row) / 2, f, k + 1);
    for (auto& a : A) a *= 2;
    K = K - (K - row) / 2 * 2;
  }
  pq_min<pair<T, int>> que;
  FOR(i, N) que.emplace(g(i, A[i]), i);
  while (K) {
    --K;
    auto [x, i] = POP(que);
    A[i]++;
    if (K) que.emplace(g(i, A[i]), i);
  }
  return A;
}

#line 1 "convex/nth_element_from_sorted_lists.hpp"

// 行の要素数が S[0], S[1], ..., S[N-1]
// O(N(logN+logK)) time
// query num: O(NlogK), あまり最適じゃないと思う
// 通ったがあやしいっぽい: https://codeforces.com/contest/1275/problem/F
// 列方向もソートされているなら matrix を使おう!!
// https://codeforces.com/problemset/problem/1034/D
template <typename T, typename F>
vi nth_element_from_sorted_lists(vi S, ll K, F f, int k = 0) {
  ll N = len(S);
  ll sm = 0;
  for (auto& x : S) sm += x >> k;
  assert(0 <= K && K <= sm);
  if (K == 0) return vi(N, 0);
  if (K == sm) return S;

  ll row = 0;
  for (auto& x : S) row += (x >= (1LL << k));

  auto g = [&](int i, ll j) -> T {
    j = ((j + 1) << k) - 1;
    return (j >= S[i] ? infty<T> : f(i, j));
  };
  vi A(N);
  if (K > row) {
    A = nth_element_from_sorted_lists<T>(S, (K - row) / 2, f, k + 1);
    for (auto& a : A) a *= 2;
    K = K - (K - row) / 2 * 2;
  }
  pq_min<pair<T, int>> que;
  FOR(i, N) que.emplace(g(i, A[i]), i);
  while (K) {
    --K;
    auto [x, i] = POP(que);
    A[i]++;
    if (K) que.emplace(g(i, A[i]), i);
  }
  return A;
}

Back to top page