Range Kth Smallest
(Src/Sequence/RangeKthSmallest.hpp)

• View this file on GitHub
• Last update: 2025-03-04 23:23:46+09:00
• Include: #include "Src/Sequence/RangeKthSmallest.hpp"

概要

長さ $N$ 静的な列に対して、「区間 $[l, r)$ の中で $k$ 番目に小さい値は？」という $Q$ 個のクエリをオフラインで計算する。

ライブラリの説明

template <class T, class U>
std::vector<T> RangeKthSmallest(const std::vector<T>& A, const std::vector<U>& Q)

A

要件

• 非空であること
• operator<が定義されていること
  • 当然だが、 $a\lt b$ であるとき、またそのときに限りa<bがtrueでないといけない

Q

要件

• Uはpublicメンバ変数l, r, kが定義されているclass, structであること
• さらに、l, r, kは符号無し整数であることが期待される。符号あり整数だとWarningが発生する
• $l\lt r \le N$ かつ $k\lt r - l$
  • l, rは0-indexedの非空な半開区間
  • kは0-indexedで、 $r - l$ 以上の値は許されない

Uは基本的には以下のstructをコピれば問題無い

struct query {
    usize l, r, k;
};

アルゴリズム

一行で: 並列二分探索 + 平面走査

$A$ をソートして重複を消した列を $(B_{1}, B_{2}, \dots, B_{M})$ とする。本問題は「区間 $[l, r)$ の中に $B_{1}, B_{2}, \dots, B_{i}$ が合計 $k$ 個以上含まれているか？」という判定問題で二分探索すれば解ける。

区間 $[1, p)$ に含まれる $B_{1}, B_{2}, \dots, B_{i}$ の数の合計を $\text{pref}(p, i)$ とする。興味のある値は $-\text{pref}(l, i) + \text{pref}(r, i)$ である。

クエリを先読みして $l, r$ を集計し、平面走査を行うことで各クエリに対して判定問題を一回ずつ解くことができる。 $\text{pref}(p, i)$ の計算にセグメント木を用いるとlog2つでできる。

セグメント木上の二分探索に想いを馳せるとさらにlogを一個落とすことができる。本ライブラリでは時間計算量 $\Theta ((N + Q)\log N)$ 、空間計算量 $\Theta (N + Q)$ を達成している。

Depends on

• 座標圧縮 (Src/Sequence/CompressedSequence.hpp)
• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Verified with

• Test/LC/range_kth_smallest.test.cpp

Code

#pragma once

#include "./CompressedSequence.hpp"
#include "../Template/TypeAlias.hpp"

#include <algorithm>
#include <bit>
#include <cassert>
#include <utility>
#include <vector>

namespace zawa {

template <class T, class U>
std::vector<T> RangeKthSmallest(const std::vector<T>& A, const std::vector<U>& Q) {
    assert(A.size());
    CompressedSequence comp{A};
    std::vector<std::vector<std::pair<usize, bool>>> event(A.size() + 1);
    for (usize i{} ; i < Q.size() ; i++) {
        const U& q{Q[i]};
        assert(q.l < q.r and q.r <= A.size());
        assert(q.k < q.r - q.l);
        event[q.l].emplace_back(i, true);
        event[q.r].emplace_back(i, false);
    }
    std::vector<usize> x(Q.size()), cnt(Q.size());
    const usize LOG{(usize)(std::bit_width(comp.size()) -
            (std::has_single_bit(comp.size()) ? 1 : 0))};
    for (usize i{LOG} ; i-- ; ) {
        std::vector<i32> seg(1 << (LOG - i));
        std::vector<i32> add(Q.size());
        for (usize j{1} ; j <= A.size() ; j++) {
            seg[comp.map(j - 1) >> i]++;
            for (auto [k, s] : event[j]) {
                add[k] += seg[(x[k] >> i)] * (s ? -1 : 1);
            }
        }
        for (u32 j{} ; j < Q.size() ; j++) {
            if (cnt[j] + add[j] <= Q[j].k) {
                x[j] += 1 << i;
                cnt[j] += add[j];
            }
        }
    } 
    std::vector<T> res(Q.size());
    for (usize i{} ; i < Q.size() ; i++) {
        res[i] = comp.inverse(x[i]);
    }
    return res;
}

} // namespace zawa

#line 2 "Src/Sequence/RangeKthSmallest.hpp"

#line 2 "Src/Sequence/CompressedSequence.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/Sequence/CompressedSequence.hpp"

#include <vector>
#include <algorithm>
#include <cassert>
#include <iterator>
#include <limits>

namespace zawa {

template <class T>
class CompressedSequence {
public:

    static constexpr u32 NotFound = std::numeric_limits<u32>::max();

    CompressedSequence() = default;

    template <class InputIterator>
    CompressedSequence(InputIterator first, InputIterator last) : comped_(first, last), f_{} {
        std::sort(comped_.begin(), comped_.end());
        comped_.erase(std::unique(comped_.begin(), comped_.end()), comped_.end());
        comped_.shrink_to_fit();
        f_.reserve(std::distance(first, last));
        for (auto it{first} ; it != last ; it++) {
            f_.emplace_back(std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), *it)));
        }
    }

    CompressedSequence(const std::vector<T>& A) : CompressedSequence(A.begin(), A.end()) {}

    inline usize size() const noexcept {
        return comped_.size();
    }

    u32 operator[](const T& v) const {
        return std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), v));
    }

    u32 upper_bound(const T& v) const {
        return std::distance(comped_.begin(), std::upper_bound(comped_.begin(), comped_.end(), v));
    }

    u32 find(const T& v) const {
        u32 i = std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), v));
        return i == comped_.size() or comped_[i] != v ? NotFound : i;
    }

    bool contains(const T& v) const {
        u32 i = std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), v));
        return i < comped_.size() and comped_[i] == v;
    }

    u32 at(const T& v) const {
        u32 res = find(v);
        assert(res != NotFound);
        return res;
    }

    inline u32 map(u32 i) const noexcept {
        assert(i < f_.size());
        return f_[i];
    }

    inline T inverse(u32 i) const noexcept {
        assert(i < size());
        return comped_[i];
    }

    inline std::vector<T> comped() const noexcept {
        return comped_;
    }

private:

    std::vector<T> comped_;

    std::vector<u32> f_;

};

} // namespace zawa
#line 5 "Src/Sequence/RangeKthSmallest.hpp"

#line 7 "Src/Sequence/RangeKthSmallest.hpp"
#include <bit>
#line 9 "Src/Sequence/RangeKthSmallest.hpp"
#include <utility>
#line 11 "Src/Sequence/RangeKthSmallest.hpp"

namespace zawa {

template <class T, class U>
std::vector<T> RangeKthSmallest(const std::vector<T>& A, const std::vector<U>& Q) {
    assert(A.size());
    CompressedSequence comp{A};
    std::vector<std::vector<std::pair<usize, bool>>> event(A.size() + 1);
    for (usize i{} ; i < Q.size() ; i++) {
        const U& q{Q[i]};
        assert(q.l < q.r and q.r <= A.size());
        assert(q.k < q.r - q.l);
        event[q.l].emplace_back(i, true);
        event[q.r].emplace_back(i, false);
    }
    std::vector<usize> x(Q.size()), cnt(Q.size());
    const usize LOG{(usize)(std::bit_width(comp.size()) -
            (std::has_single_bit(comp.size()) ? 1 : 0))};
    for (usize i{LOG} ; i-- ; ) {
        std::vector<i32> seg(1 << (LOG - i));
        std::vector<i32> add(Q.size());
        for (usize j{1} ; j <= A.size() ; j++) {
            seg[comp.map(j - 1) >> i]++;
            for (auto [k, s] : event[j]) {
                add[k] += seg[(x[k] >> i)] * (s ? -1 : 1);
            }
        }
        for (u32 j{} ; j < Q.size() ; j++) {
            if (cnt[j] + add[j] <= Q[j].k) {
                x[j] += 1 << i;
                cnt[j] += add[j];
            }
        }
    } 
    std::vector<T> res(Q.size());
    for (usize i{} ; i < Q.size() ; i++) {
        res[i] = comp.inverse(x[i]);
    }
    return res;
}

} // namespace zawa

Back to top page