src/verify/yosupo-ordered-set.test.cpp

• View this file on GitHub
• Last update: 2026-01-18 08:41:56+09:00
• Problem: https://judge.yosupo.jp/problem/ordered_set

Depends on

• Fenwick Tree (src/data_structure/fenwicktree.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/ordered_set"
#include "../../cp-library/fenwicktree"

#include <bits/stdc++.h>
using namespace std;

template <class T>
struct Tree {
    const vector<T> xs;
    akTARDIGRADE13::fenwick_tree<int> ft;
    const int n;

    Tree() : n(0) {}
    explicit Tree(const vector<T>& _xs) : xs(_xs), n(xs.size()), ft(xs.size()) {}

    int idx_of(const T& x) const { return distance(xs.begin(), ranges::lower_bound(xs, x)); }

    bool contains(const T& x) const {
        int i = idx_of(x);
        if (i >= n || xs[i] != x) return false;
        return ft.sum(i, i + 1) != 0;
    }

    bool insert(const T& x) {
        int i = idx_of(x);
        if (i >= n || xs[i] != x) return false;
        if (ft.sum(i, i + 1) != 0) return false;
        ft.add(i, 1);
        return true;
    }

    bool erase(const T& x) {
        int i = idx_of(x);
        if (i >= n || xs[i] != x) return false;
        if (ft.sum(i, i + 1) == 0) return false;
        ft.add(i, -1);
        return true;
    }

    int size() const { return ft.sum(0, n); }

    // # of elements < x
    int count_lt(const T& x) const {
        int r = distance(xs.begin(), ranges::lower_bound(xs, x));
        return ft.sum(0, r);
    }

    // # of elements <= x
    int count_le(const T& x) const {
        int r = distance(xs.begin(), ranges::upper_bound(xs, x));
        return ft.sum(0, r);
    }

    // k-th smallest (0-indexed). none if out of range.
    optional<T> kth(int k) const {
        int tot = ft.sum(0, n);
        if (k < 0 || k >= tot) return nullopt;
        size_t len = ft.lower_bound(k + 1);
        int idx = (int)len - 1;
        return xs[idx];
    }

    // min element >= x
    optional<T> next_ge(const T& x) const {
        int pos = distance(xs.begin(), ranges::lower_bound(xs, x));
        int k = ft.sum(0, pos);
        return kth(k);
    }

    // max element <= x
    optional<T> prev_le(const T& x) const {
        int cnt = count_le(x);
        if (cnt == 0) return nullopt;
        return kth(cnt - 1);
    }
};

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<int> a(N), t(Q), x(Q), num;
    num.reserve(N + Q);
    for (int i = 0; i < N; ++i) {
        cin >> a[i];
        num.push_back(a[i]);
    }
    for (int i = 0; i < Q; ++i) {
        cin >> t[i] >> x[i];
        num.push_back(x[i]);
    }
    ranges::sort(num);
    num.erase(ranges::unique(num).begin(), num.end());
    Tree<int> tree(num);

    for (int i = 0; i < N; ++i) {
        tree.insert(a[i]);
    }

    for (int i = 0; i < Q; ++i) {
        if (t[i] == 0) {
            tree.insert(x[i]);
        } else if (t[i] == 1) {
            tree.erase(x[i]);
        } else if (t[i] == 2) {
            auto ans = tree.kth(x[i] - 1);
            cout << (ans ? *ans : -1) << '\n';
        } else if (t[i] == 3) {
            cout << tree.count_le(x[i]) << '\n';
        } else if (t[i] == 4) {
            auto ans = tree.prev_le(x[i]);
            cout << (ans ? *ans : -1) << '\n';
        } else {
            auto ans = tree.next_ge(x[i]);
            cout << (ans ? *ans : -1) << '\n';
        }
    }
    return 0;
}

#line 1 "src/verify/yosupo-ordered-set.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/ordered_set"
#line 2 "src/data_structure/fenwicktree.hpp"

#include <cassert>
#include <cstddef>
#include <vector>

namespace akTARDIGRADE13 {

/**
 * @brief Fenwick Tree
 *
 * 点更新および区間和クエリを高速に処理するためのデータ構造。
 *
 * - 外部から見えるインデックスは **0-indexed**
 * - 内部では **1-indexed** の Fenwick Tree を使用
 * - 半開区間 `[l, r)` を基本とする
 *
 * @tparam T
 * 加算・比較が可能な型。
 *
 * @note
 * `lower_bound` / `upper_bound` を使用する場合、
 * 木に格納される値は **非負**である必要がある。
 */
template <class T>
struct fenwick_tree {

public:
    /**
     * @brief サイズ 0 の Fenwick Tree を構築する。
     *
     * @complexity O(1)
     */
    fenwick_tree() : _n(0) {}

    /**
     * @brief サイズ @p n の Fenwick Tree を構築する（全要素 0 初期化）。
     *
     * @param n 配列サイズ
     *
     * @complexity O(n)
     */
    explicit fenwick_tree(int n) : _n(n), data(n) {}

    /**
     * @brief 位置 @p p に値 @p x を加算する。
     *
     * @param p 加算する位置（0-indexed）
     * @param x 加算する値
     *
     * @pre 0 <= p < n
     *
     * @complexity O(log n)
     */
    void add(int p, T x) {
        assert(0 <= p && p < _n);
        p++;
        while (p <= _n) {
            data[p - 1] += x;
            p += p & -p;
        }
    }

    /**
     * @brief 区間 `[l, r)` の総和を返す。
     *
     * @param l 左端（含む）
     * @param r 右端（含まない）
     *
     * @return 区間 `[l, r)` に含まれる要素の総和
     *
     * @pre 0 <= l <= r <= n
     *
     * @complexity O(log n)
     */
    T sum(int l, int r) const {
        assert(0 <= l && l <= r && r <= _n);
        return sum(r) - sum(l);
    }

    /**
     * @brief prefix sum が @p w 以上となる最小の index を返す。
     *
     * 正確には、次を満たす最小の @p i を返す：
     *
     * ```
     * sum([0, i)) >= w
     * ```
     *
     * 返り値は **要素の index ではなく prefix の長さ**である点に注意。
     *
     * @param w 目標とする累積和
     *
     * @return
     * - `w <= 0` の場合：`0`
     * - 条件を満たす最小の `i`
     *
     * @note
     * - この関数は格納される値が **非負**であることを仮定する。
     * - 条件を満たす `i` が存在しない場合、`n+1` が返る。
     *
     * @complexity O(log n)
     */
    std::size_t lower_bound(T w) const {
        if (w <= T{}) return 0;
        std::size_t idx = 0;
        T cur = T{};
        int step = 1;
        while (step < _n) step <<= 1;
        while (step) {
            std::size_t nxt = idx + step;
            if (nxt <= _n) {
                T val = cur + data[nxt - 1];
                if (val < w) {
                    idx = nxt;
                    cur = val;
                }
            }
            step >>= 1;
        }
        return idx + 1;
    }

    /**
     * @brief prefix sum が @p w を超える最小の index を返す。
     *
     * 正確には、次を満たす最小の @p i を返す：
     *
     * ```
     * sum([0, i)) > w
     * ```
     *
     * 返り値は **要素の index ではなく prefix の長さ**である。
     *
     * @param w 目標とする累積和
     *
     * @return
     * - `w <= 0` の場合：`0`
     * - 条件を満たす最小の `i`
     *
     * @note
     * - この関数は格納される値が **非負**であることを仮定する。
     * - 条件を満たす `i` が存在しない場合、`n+1` が返る。
     *
     * @complexity O(log n)
     */
    std::size_t upper_bound(T w) const {
        if (w <= T{}) return 0;
        std::size_t idx = 0;
        T cur = T{};
        int step = 1;
        while (step < _n) step <<= 1;
        while (step) {
            std::size_t nxt = idx + step;
            if (nxt <= _n) {
                T val = cur + data[nxt - 1];
                if (val <= w) {
                    idx = nxt;
                    cur = val;
                }
            }
            step >>= 1;
        }
        return idx + 1;
    }

private:
    int _n;              ///< 配列サイズ
    std::vector<T> data; ///< Fenwick Tree の内部配列（1-indexed 相当）

    /**
     * @brief prefix 区間 `[0, r)` の総和を返す。
     *
     * @param r 右端（含まない）
     *
     * @return 区間 `[0, r)` の総和
     *
     * @complexity O(log n)
     */
    T sum(int r) const {
        T s = 0;
        while (r > 0) {
            s += data[r - 1];
            r -= r & -r;
        }
        return s;
    }
};

} // namespace akTARDIGRADE13
#line 3 "src/verify/yosupo-ordered-set.test.cpp"

#include <bits/stdc++.h>
using namespace std;

template <class T>
struct Tree {
    const vector<T> xs;
    akTARDIGRADE13::fenwick_tree<int> ft;
    const int n;

    Tree() : n(0) {}
    explicit Tree(const vector<T>& _xs) : xs(_xs), n(xs.size()), ft(xs.size()) {}

    int idx_of(const T& x) const { return distance(xs.begin(), ranges::lower_bound(xs, x)); }

    bool contains(const T& x) const {
        int i = idx_of(x);
        if (i >= n || xs[i] != x) return false;
        return ft.sum(i, i + 1) != 0;
    }

    bool insert(const T& x) {
        int i = idx_of(x);
        if (i >= n || xs[i] != x) return false;
        if (ft.sum(i, i + 1) != 0) return false;
        ft.add(i, 1);
        return true;
    }

    bool erase(const T& x) {
        int i = idx_of(x);
        if (i >= n || xs[i] != x) return false;
        if (ft.sum(i, i + 1) == 0) return false;
        ft.add(i, -1);
        return true;
    }

    int size() const { return ft.sum(0, n); }

    // # of elements < x
    int count_lt(const T& x) const {
        int r = distance(xs.begin(), ranges::lower_bound(xs, x));
        return ft.sum(0, r);
    }

    // # of elements <= x
    int count_le(const T& x) const {
        int r = distance(xs.begin(), ranges::upper_bound(xs, x));
        return ft.sum(0, r);
    }

    // k-th smallest (0-indexed). none if out of range.
    optional<T> kth(int k) const {
        int tot = ft.sum(0, n);
        if (k < 0 || k >= tot) return nullopt;
        size_t len = ft.lower_bound(k + 1);
        int idx = (int)len - 1;
        return xs[idx];
    }

    // min element >= x
    optional<T> next_ge(const T& x) const {
        int pos = distance(xs.begin(), ranges::lower_bound(xs, x));
        int k = ft.sum(0, pos);
        return kth(k);
    }

    // max element <= x
    optional<T> prev_le(const T& x) const {
        int cnt = count_le(x);
        if (cnt == 0) return nullopt;
        return kth(cnt - 1);
    }
};

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<int> a(N), t(Q), x(Q), num;
    num.reserve(N + Q);
    for (int i = 0; i < N; ++i) {
        cin >> a[i];
        num.push_back(a[i]);
    }
    for (int i = 0; i < Q; ++i) {
        cin >> t[i] >> x[i];
        num.push_back(x[i]);
    }
    ranges::sort(num);
    num.erase(ranges::unique(num).begin(), num.end());
    Tree<int> tree(num);

    for (int i = 0; i < N; ++i) {
        tree.insert(a[i]);
    }

    for (int i = 0; i < Q; ++i) {
        if (t[i] == 0) {
            tree.insert(x[i]);
        } else if (t[i] == 1) {
            tree.erase(x[i]);
        } else if (t[i] == 2) {
            auto ans = tree.kth(x[i] - 1);
            cout << (ans ? *ans : -1) << '\n';
        } else if (t[i] == 3) {
            cout << tree.count_le(x[i]) << '\n';
        } else if (t[i] == 4) {
            auto ans = tree.prev_le(x[i]);
            cout << (ans ? *ans : -1) << '\n';
        } else {
            auto ans = tree.next_ge(x[i]);
            cout << (ans ? *ans : -1) << '\n';
        }
    }
    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	58 ms	18 MB
g++	example_01	AC	12 ms	18 MB
g++	max_random_00	AC	453 ms	26 MB
g++	max_random_01	AC	465 ms	26 MB
g++	max_random_02	AC	377 ms	25 MB
g++	max_random_03	AC	377 ms	25 MB
g++	max_random_04	AC	329 ms	25 MB
g++	max_random_05	AC	467 ms	27 MB
g++	max_random_06	AC	512 ms	28 MB
g++	max_random_07	AC	532 ms	28 MB
g++	max_random_08	AC	474 ms	26 MB
g++	max_random_09	AC	479 ms	28 MB
g++	max_random_10	AC	382 ms	25 MB
g++	max_random_11	AC	379 ms	25 MB
g++	max_random_12	AC	339 ms	25 MB
g++	max_random_13	AC	489 ms	28 MB
g++	max_random_14	AC	561 ms	28 MB
g++	max_random_15	AC	562 ms	27 MB
g++	max_random_16	AC	893 ms	36 MB
g++	max_random_17	AC	899 ms	37 MB
g++	max_random_18	AC	809 ms	35 MB
g++	max_random_19	AC	828 ms	35 MB
g++	max_random_20	AC	840 ms	36 MB
g++	max_random_21	AC	866 ms	37 MB
g++	max_random_22	AC	954 ms	37 MB
g++	max_random_23	AC	959 ms	37 MB
g++	max_random_24	AC	906 ms	36 MB
g++	max_random_25	AC	942 ms	37 MB
g++	max_random_26	AC	814 ms	35 MB
g++	max_random_27	AC	804 ms	35 MB
g++	max_random_28	AC	836 ms	36 MB
g++	max_random_29	AC	870 ms	36 MB
g++	max_random_30	AC	980 ms	36 MB
g++	max_random_31	AC	1015 ms	37 MB
g++	small_00	AC	261 ms	24 MB
g++	small_01	AC	267 ms	25 MB
g++	small_02	AC	267 ms	25 MB

Back to top page