tests/wavelet_matrix/kth_smallest.test.cpp

View this file on GitHub
Last update: 2025-02-10 20:05:30+09:00
Problem: https://judge.yosupo.jp/problem/range_kth_smallest

Depends on

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/range_kth_smallest

#include <iostream>
#include "tools/wavelet_matrix.hpp"

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N, Q;
  std::cin >> N >> Q;
  tools::wavelet_matrix<int> wm;
  for (int i = 0; i < N; ++i) {
    int a_i;
    std::cin >> a_i;
    wm.add_point(i, a_i);
  }
  wm.build();

  for (int q = 0; q < Q; ++q) {
    int l, r, k;
    std::cin >> l >> r >> k;
    std::cout << wm.get_point(wm.kth_smallest(l, r, k)).second << "\n";
  }

  return 0;
}

#line 1 "tests/wavelet_matrix/kth_smallest.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/range_kth_smallest

#include <iostream>
#line 1 "tools/wavelet_matrix.hpp"



#include <vector>
#include <utility>
#include <cstddef>
#include <cassert>
#include <algorithm>
#include <iterator>
#include <array>
#include <tuple>
#include <optional>
#line 1 "tools/bit_vector.hpp"



#include <cstdint>
#line 6 "tools/bit_vector.hpp"
#include <immintrin.h>

// Source: https://nyaannyaan.github.io/library/data-structure-2d/wavelet-matrix.hpp.html
// License: CC0 1.0 Universal
// Author: Nyaan

namespace tools {
  class bit_vector {
  private:
    using u32 = ::std::uint32_t;
    using i64 = ::std::int64_t;
    using u64 = ::std::uint64_t;

    static constexpr u32 w = 64;
    ::std::vector<u64> m_block;
    ::std::vector<u32> m_count;
    u32 m_size, m_zeros;

  public:
    bit_vector() {}
    explicit bit_vector(int _n) { init(_n); }
    __attribute__((optimize("O3", "unroll-loops"))) void init(int _n) {
      m_size = m_zeros = _n;
      m_block.resize(m_size / w + 1, 0);
      m_count.resize(m_block.size(), 0);
    }

    u32 size() const { return m_size; }
    inline u32 get(u32 i) const { return u32(m_block[i / w] >> (i % w)) & 1u; }
    inline void set(u32 i) { m_block[i / w] |= 1LL << (i % w); }

    __attribute__((target("popcnt"))) void build() {
      for (u32 i = 1; i < m_block.size(); ++i)
        m_count[i] = m_count[i - 1] + _mm_popcnt_u64(m_block[i - 1]);
      m_zeros = rank0(m_size);
    }

    u32 zeros() const { return m_zeros; }
    inline u32 rank0(u32 i) const { return i - rank1(i); }
    __attribute__((target("bmi2,popcnt"))) inline u32 rank1(u32 i) const {
      return m_count[i / w] + _mm_popcnt_u64(_bzhi_u64(m_block[i / w], i % w));
    }
  };
}


#line 1 "tools/lower_bound.hpp"



#line 6 "tools/lower_bound.hpp"

namespace tools {

  template <class ForwardIterator, class T>
  auto lower_bound(ForwardIterator first, ForwardIterator last, const T& value) {
    return ::std::distance(first, ::std::lower_bound(first, last, value));
  }

  template <class ForwardIterator, class T, class Compare>
  auto lower_bound(ForwardIterator first, ForwardIterator last, const T& value, Compare comp) {
    return ::std::distance(first, ::std::lower_bound(first, last, value, comp));
  }
}


#line 1 "tools/less_by_first.hpp"



#line 5 "tools/less_by_first.hpp"

namespace tools {

  class less_by_first {
  public:
    template <class T1, class T2>
    bool operator()(const ::std::pair<T1, T2>& x, const ::std::pair<T1, T2>& y) const {
      return x.first < y.first;
    }
  };
}


#line 1 "tools/floor_log2.hpp"



#line 1 "tools/bit_width.hpp"



#include <bit>
#line 6 "tools/bit_width.hpp"
#include <type_traits>
#line 1 "tools/is_integral.hpp"



#line 5 "tools/is_integral.hpp"

namespace tools {
  template <typename T>
  struct is_integral : ::std::is_integral<T> {};

  template <typename T>
  inline constexpr bool is_integral_v = ::tools::is_integral<T>::value;
}


#line 1 "tools/is_signed.hpp"



#line 5 "tools/is_signed.hpp"

namespace tools {
  template <typename T>
  struct is_signed : ::std::is_signed<T> {};

  template <typename T>
  inline constexpr bool is_signed_v = ::tools::is_signed<T>::value;
}


#line 1 "tools/make_unsigned.hpp"



#line 5 "tools/make_unsigned.hpp"

namespace tools {
  template <typename T>
  struct make_unsigned : ::std::make_unsigned<T> {};

  template <typename T>
  using make_unsigned_t = typename ::tools::make_unsigned<T>::type;
}


#line 10 "tools/bit_width.hpp"

namespace tools {
  template <typename T>
  constexpr int bit_width(T) noexcept;

  template <typename T>
  constexpr int bit_width(const T x) noexcept {
    static_assert(::tools::is_integral_v<T> && !::std::is_same_v<::std::remove_cv_t<T>, bool>);
    if constexpr (::tools::is_signed_v<T>) {
      assert(x >= 0);
      return ::tools::bit_width<::tools::make_unsigned_t<T>>(x);
    } else {
      return ::std::bit_width(x);
    }
  }
}


#line 6 "tools/floor_log2.hpp"

namespace tools {
  template <typename T>
  constexpr T floor_log2(T x) noexcept {
    assert(x > 0);
    return ::tools::bit_width(x) - 1;
  }
}


#line 17 "tools/wavelet_matrix.hpp"

namespace tools {
  template <typename T>
  class wavelet_matrix {
  private:
    ::std::vector<::std::pair<T, T>> m_ps;
    ::std::vector<::std::pair<T, ::std::size_t>> m_xs;
    ::std::vector<T> m_ys;
    ::std::vector<::std::size_t> m_is;
    ::std::vector<::tools::bit_vector> m_bvs;

    ::std::size_t iid(const ::std::size_t i) const {
      return ::tools::lower_bound(this->m_xs.begin(), this->m_xs.end(), ::std::make_pair(this->m_ps[i].first, i));
    }
    ::std::size_t xid(const T& x) const {
      return ::tools::lower_bound(this->m_xs.begin(), this->m_xs.end(), ::std::make_pair(x, ::std::size_t{}), ::tools::less_by_first());
    }
    ::std::size_t yid(const T& y) const {
      return ::tools::lower_bound(this->m_ys.begin(), this->m_ys.end(), y);
    }
    ::std::size_t lg() const {
      return this->m_bvs.size();
    }
    bool built() const {
      return this->lg() > 0;
    }

  public:
    wavelet_matrix() = default;
    wavelet_matrix(const ::tools::wavelet_matrix<T>&) = default;
    wavelet_matrix(::tools::wavelet_matrix<T>&&) = default;
    ~wavelet_matrix() = default;
    ::tools::wavelet_matrix<T>& operator=(const ::tools::wavelet_matrix<T>&) = default;
    ::tools::wavelet_matrix<T>& operator=(::tools::wavelet_matrix<T>&&) = default;

    ::std::size_t size() const {
      return this->m_ps.size();
    }

    ::std::size_t add_point(const T& x, const T& y) {
      assert(!this->built());
      this->m_ps.emplace_back(x, y);
      return this->size() - 1;
    }

    ::std::pair<T, T> get_point(const ::std::size_t i) const {
      assert(i < this->size());
      return this->m_ps[i];
    }

    const ::std::vector<::std::pair<T, T>>& points() const {
      return this->m_ps;
    }

    ::std::vector<::std::vector<::std::size_t>> build() {
      assert(!this->built());

      this->m_xs.reserve(this->size());
      for (::std::size_t i = 0; i < this->size(); ++i) {
        this->m_xs.emplace_back(this->m_ps[i].first, i);
      }
      ::std::sort(this->m_xs.begin(), this->m_xs.end());

      this->m_ys.reserve(this->size());
      ::std::transform(this->m_ps.begin(), this->m_ps.end(), ::std::back_inserter(this->m_ys), [](const auto& p) { return p.second; });
      ::std::sort(this->m_ys.begin(), this->m_ys.end());
      this->m_ys.erase(::std::unique(this->m_ys.begin(), this->m_ys.end()), this->m_ys.end());

      const auto n = ::std::max<::std::size_t>(this->size(), 1);
      this->m_bvs.assign(::tools::floor_log2(::std::max<::std::size_t>(this->m_ys.size(), 1)) + 1, ::tools::bit_vector(n));

      ::std::vector<::std::size_t> cur;
      cur.reserve(n);
      ::std::transform(this->m_xs.begin(), this->m_xs.end(), ::std::back_inserter(cur), [&](const auto& p) { return this->yid(this->m_ps[p.second].second); });
      cur.resize(n);
      ::std::vector<::std::size_t> nxt(n);

      auto res = ::std::vector(this->lg() + 1, ::std::vector<::std::size_t>(n));
      ::std::transform(this->m_xs.begin(), this->m_xs.end(), res.back().begin(), [&](const auto& p) { return p.second; });

      for (::std::size_t h = this->lg(); h --> 0;) {
        for (::std::size_t i = 0; i < n; ++i) {
          if ((cur[i] >> h) & 1) {
            this->m_bvs[h].set(i);
          }
        }
        this->m_bvs[h].build();
        ::std::array<::std::size_t, 2> offsets = {0, this->m_bvs[h].zeros()};
        for (::std::size_t i = 0; i < n; ++i) {
          auto& offset = offsets[this->m_bvs[h].get(i)];
          nxt[offset] = cur[i];
          res[h][offset] = res[h + 1][i];
          ++offset;
        }
        ::std::swap(cur, nxt);
      }

      this->m_is = res.front();
      res.pop_back();
      return res;
    }

    ::std::vector<::std::pair<::std::size_t, ::std::size_t>> apply(::std::size_t i) const {
      assert(this->built());
      assert(i < this->size());

      ::std::vector<::std::pair<::std::size_t, ::std::size_t>> res(this->lg());
      i = this->iid(i);
      for (::std::size_t h = this->lg(); h --> 0;) {
        i = this->m_bvs[h].get(i) ? this->m_bvs[h].zeros() + this->m_bvs[h].rank1(i) : this->m_bvs[h].rank0(i);
        res[h] = ::std::make_pair(h, i);
      }
      return res;
    }

    ::std::size_t kth_smallest(const T& l, const T& r, ::std::size_t k) const {
      assert(this->built());
      assert(l <= r);

      auto lid = this->xid(l);
      auto rid = this->xid(r);

      assert(k < rid - lid);

      for (::std::size_t h = this->lg(); h --> 0;) {
        const auto l0 = this->m_bvs[h].rank0(lid);
        const auto r0 = this->m_bvs[h].rank0(rid);
        if (k < r0 - l0) {
          lid = l0;
          rid = r0;
        } else {
          k -= r0 - l0;
          lid += this->m_bvs[h].zeros() - l0;
          rid += this->m_bvs[h].zeros() - r0;
        }
      }

      return this->m_is[lid + k];
    }

    ::std::size_t kth_largest(const T& l, const T& r, const ::std::size_t k) const {
      assert(this->built());
      assert(l <= r);

      const auto lid = this->xid(l);
      const auto rid = this->xid(r);

      assert(k < rid - lid);

      return this->kth_smallest(l, r, rid - lid - k - 1);
    }

    ::std::vector<::std::tuple<::std::size_t, ::std::size_t, ::std::size_t>> range_prod(const T& l, const T& r, const T& u) const {
      assert(this->built());
      assert(l <= r);

      auto lid = this->xid(l);
      auto rid = this->xid(r);
      const auto uid = this->yid(u);

      ::std::vector<::std::tuple<::std::size_t, ::std::size_t, ::std::size_t>> res(this->lg());
      for (::std::size_t h = this->lg(); h --> 0;) {
        const auto l0 = this->m_bvs[h].rank0(lid);
        const auto r0 = this->m_bvs[h].rank0(rid);
        if ((uid >> h) & 1) {
          res[h] = ::std::make_tuple(h, l0, r0);
          lid += this->m_bvs[h].zeros() - l0;
          rid += this->m_bvs[h].zeros() - r0;
        } else {
          res[h] = ::std::make_tuple(h, 0, 0);
          lid = l0;
          rid = r0;
        }
      }
      return res;
    }

    ::std::size_t range_freq(const T& l, const T& r) const {
      assert(this->built());
      assert(l <= r);

      return this->xid(r) - this->xid(l);
    }

    ::std::size_t range_freq(const T& l, const T& r, const T& u) const {
      assert(this->built());
      assert(l <= r);

      auto lid = this->xid(l);
      auto rid = this->xid(r);
      const auto uid = this->yid(u);

      ::std::size_t res = 0;
      for (::std::size_t h = this->lg(); h --> 0;) {
        const auto l0 = this->m_bvs[h].rank0(lid);
        const auto r0 = this->m_bvs[h].rank0(rid);
        if ((uid >> h) & 1) {
          res += r0 - l0;
          lid += this->m_bvs[h].zeros() - l0;
          rid += this->m_bvs[h].zeros() - r0;
        } else {
          lid = l0;
          rid = r0;
        }
      }
      return res;
    }

    ::std::size_t range_freq(const T& l, const T& r, const T& d, const T& u) const {
      assert(this->built());
      assert(l <= r);
      assert(d <= u);

      return this->range_freq(l, r, u) - this->range_freq(l, r, d);
    }

    ::std::optional<T> prev_value(const T& l, const T& r, const T& u) const {
      assert(this->built());
      assert(l <= r);

      const auto k = this->range_freq(l, r, u);
      return k == 0 ? ::std::nullopt : ::std::make_optional(this->m_ps[this->kth_smallest(l, r, k - 1)].second);
    }

    ::std::optional<T> next_value(const T& l, const T& r, const T& d) const {
      assert(this->built());
      assert(l <= r);

      const auto k = this->range_freq(l, r, d);
      return k == this->range_freq(l, r) ? ::std::nullopt : ::std::make_optional(this->m_ps[this->kth_smallest(l, r, k)].second);
    }

    ::std::pair<typename ::std::vector<::std::size_t>::const_iterator, typename ::std::vector<::std::size_t>::const_iterator> prev_points(const T& l, const T& r, const T& u) const {
      assert(this->built());
      assert(l <= r);

      auto lid = this->xid(l);
      auto rid = this->xid(r);
      auto k = this->range_freq(l, r, u);
      if (k == 0) return ::std::make_pair(this->m_is.cend(), this->m_is.cend());
      --k;

      for (::std::size_t h = this->lg(); h --> 0;) {
        const auto l0 = this->m_bvs[h].rank0(lid);
        const auto r0 = this->m_bvs[h].rank0(rid);
        if (k < r0 - l0) {
          lid = l0;
          rid = r0;
        } else {
          k -= r0 - l0;
          lid += this->m_bvs[h].zeros() - l0;
          rid += this->m_bvs[h].zeros() - r0;
        }
      }

      return ::std::make_pair(this->m_is.cbegin() + lid, this->m_is.cbegin() + rid);
    }

    ::std::pair<typename ::std::vector<::std::size_t>::const_iterator, typename ::std::vector<::std::size_t>::const_iterator> next_points(const T& l, const T& r, const T& d) const {
      assert(this->built());
      assert(l <= r);

      auto lid = this->xid(l);
      auto rid = this->xid(r);
      auto k = this->range_freq(l, r, d);
      if (k == rid - lid) return ::std::make_pair(this->m_is.cend(), this->m_is.cend());

      for (::std::size_t h = this->lg(); h --> 0;) {
        const auto l0 = this->m_bvs[h].rank0(lid);
        const auto r0 = this->m_bvs[h].rank0(rid);
        if (k < r0 - l0) {
          lid = l0;
          rid = r0;
        } else {
          k -= r0 - l0;
          lid += this->m_bvs[h].zeros() - l0;
          rid += this->m_bvs[h].zeros() - r0;
        }
      }

      return ::std::make_pair(this->m_is.cbegin() + lid, this->m_is.cbegin() + rid);
    }
  };
}


#line 5 "tests/wavelet_matrix/kth_smallest.test.cpp"

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N, Q;
  std::cin >> N >> Q;
  tools::wavelet_matrix<int> wm;
  for (int i = 0; i < N; ++i) {
    int a_i;
    std::cin >> a_i;
    wm.add_point(i, a_i);
  }
  wm.build();

  for (int q = 0; q < Q; ++q) {
    int l, r, k;
    std::cin >> l >> r >> k;
    std::cout << wm.get_point(wm.kth_smallest(l, r, k)).second << "\n";
  }

  return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	all_zero_00	AC	6 ms	4 MB
g++	dense_large_a_00	AC	51 ms	4 MB
g++	dense_small_a_00	AC	43 ms	4 MB
g++	example_00	AC	5 ms	4 MB
g++	max_random_00	AC	283 ms	44 MB
g++	max_random_01	AC	285 ms	44 MB
g++	max_random_02	AC	280 ms	44 MB
g++	max_random_03	AC	282 ms	44 MB
g++	max_random_04	AC	285 ms	44 MB
g++	random_00	AC	190 ms	28 MB
g++	random_01	AC	208 ms	34 MB
g++	random_02	AC	117 ms	13 MB
g++	random_03	AC	105 ms	37 MB
g++	random_04	AC	63 ms	6 MB
g++	small_00	AC	5 ms	4 MB
g++	small_01	AC	5 ms	4 MB
g++	small_02	AC	5 ms	4 MB
g++	small_03	AC	5 ms	4 MB
g++	small_04	AC	5 ms	4 MB
g++	small_05	AC	5 ms	4 MB
g++	small_06	AC	5 ms	4 MB
g++	small_07	AC	5 ms	4 MB
g++	small_08	AC	5 ms	4 MB
g++	small_09	AC	5 ms	4 MB