Ordered Set (verify/structures/fenwick/ordered_set.test.cpp)

• View this file on GitHub
• Last update: 2024-11-14 22:26:18+01:00
• Problem: https://judge.yosupo.jp/problem/ordered_set

Depends on

• cp-algo/structures/bit_array.hpp
• cp-algo/structures/fenwick.hpp
• cp-algo/structures/fenwick_set.hpp
• cp-algo/util/bit.hpp
• cp-algo/util/compress_coords.hpp
• cp-algo/util/sort.hpp

Code

// @brief Ordered Set
#define PROBLEM "https://judge.yosupo.jp/problem/ordered_set"
#pragma GCC optimize("Ofast,unroll-loops")
#include "cp-algo/structures/fenwick_set.hpp"
#include "cp-algo/util/compress_coords.hpp"
#include <bits/stdc++.h>

using namespace std;
using cp_algo::structures::fenwick_set;

void solve() {
    int n, q;
    cin >> n >> q;
    vector a(n, 0);
    vector<reference_wrapper<int>> coords;
    for(auto &it: a) {
        cin >> it;
        coords.push_back(ref(it));
    }
    vector queries(q, pair{0, 0});
    for(auto &[t, x]: queries) {
        cin >> t >> x;
        if(t != 2) {
            coords.push_back(ref(x));
        }
    }
    auto values = cp_algo::compress_coords(coords);
    const int maxc = 1e6;
    fenwick_set<maxc> me(a);
    for(auto [t, x]: queries) {
        if(t == 0) {
            me.insert(x);
        } else if(t == 1) {
            me.erase(x);
        } else if(t == 2) {
            int res = me.find_by_order(x-1);
            cout << (res == -1 ? -1 : values[res]) << '\n';
        } else if(t == 3) {
            cout << me.order_of_key(x+1) << '\n';
        } else if(t == 4) {
            int res = me.pre_upper_bound(x);
            cout << (res == -1 ? -1 : values[res]) << '\n';
        } else if(t == 5) {
            int res = me.lower_bound(x);
            cout << (res == -1 ? -1 : values[res]) << '\n';
        }
    }
}

signed main() {
    //freopen("input.txt", "r", stdin);
    ios::sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    while(t--) {
        solve();
    }
}

#line 1 "verify/structures/fenwick/ordered_set.test.cpp"
// @brief Ordered Set
#define PROBLEM "https://judge.yosupo.jp/problem/ordered_set"
#pragma GCC optimize("Ofast,unroll-loops")
#line 1 "cp-algo/structures/fenwick_set.hpp"


#line 1 "cp-algo/structures/fenwick.hpp"


#include <cassert>
#include <vector>
namespace cp_algo::structures {
    template<typename T, typename Container = std::vector<T>>
    struct fenwick {
        size_t n;
        Container data;

        fenwick(auto &&range) {
            assign(move(range));
        }
        void to_prefix_sums() {
            for(size_t i = 1; i < n; i++) {
                if(i + (i & -i) <= n) {
                    data[i + (i & -i)] += data[i];
                }
            }
        }
        void assign(auto &&range) {
            n = size(range) - 1;
            data = move(range);
            to_prefix_sums();
        }
        void add(size_t x, T const& v) {
            for(++x; x <= n; x += x & -x) {
                data[x] += v;
            }
        }
        // sum of [0, r)
        T prefix_sum(size_t r) const {
            assert(r <= n);
            T res = 0;
            for(; r; r -= r & -r) {
                res += data[r];
            }
            return res;
        }
        // sum of [l, r)
        T range_sum(size_t l, size_t r) const {
            return prefix_sum(r) - prefix_sum(l);
        }
        // Last x s.t. k = prefix_sum(x) + r for r > 0
        // Assumes data[x] >= 0 for all x, returns [x, r]
        auto prefix_lower_bound(T k) const {
            int x = 0;
            for(size_t i = std::bit_floor(n); i; i /= 2) {
                if(x + i <= n && data[x + i] < k) {
                    k -= data[x + i];
                    x += i;
                }
            }
            return std::pair{x, k};
        }
    };
}

#line 1 "cp-algo/structures/bit_array.hpp"


#line 1 "cp-algo/util/bit.hpp"


#include <immintrin.h>
#include <cstdint>
#include <array>
#include <bit>
namespace cp_algo {
    template<typename Uint>
    constexpr size_t bit_width = sizeof(Uint) * 8;

    size_t order_of_bit(auto x, size_t k) {
        return k ? std::popcount(x << (bit_width<decltype(x)> - k)) : 0;
    }
    // Requires GCC target("popcnt,bmi2")
    size_t kth_set_bit(uint64_t x, size_t k) {
        return std::countr_zero(_pdep_u64(1ULL << k, x));
    }
    template<int fl = 0>
    void with_bit_floor(size_t n, auto &&callback) {
        if constexpr (fl >= 63) {
            return;
        } else if (n >> (fl + 1)) {
            with_bit_floor<fl + 1>(n, callback);
        } else {
            callback.template operator()<1ULL << fl>();
        }
    }
}

#line 4 "cp-algo/structures/bit_array.hpp"
namespace cp_algo::structures {
    template<size_t N, typename Uint = uint64_t>
    struct bit_array {
        static constexpr size_t width = bit_width<Uint>;
        static constexpr size_t blocks = N / width + 1;
        std::array<Uint, blocks> data = {};

        uint64_t word(size_t x) const {
            return data[x];
        }
        void set(size_t x) {
            data[x / width] |= 1ULL << (x % width);
        }
        void flip(size_t x) {
            data[x / width] ^= 1ULL << (x % width);
        }
        bool test(size_t x) const {
            return (data[x / width] >> (x % width)) & 1;
        }
        bool operator[](size_t x) const {
            return test(x);
        }
    };
}

#line 5 "cp-algo/structures/fenwick_set.hpp"
namespace cp_algo::structures {
    template<size_t maxc, typename Uint = uint64_t>
    using popcount_array = std::array<int, maxc / bit_width<Uint> + 1>;
    // fenwick-based set for [0, maxc)
    template<size_t maxc, typename Uint = uint64_t>
    struct fenwick_set: fenwick<int, popcount_array<maxc, Uint>> {
        using Base = fenwick<int, popcount_array<maxc, Uint>>;
        static constexpr size_t word = bit_width<Uint>;
        size_t sz = 0;
        bit_array<maxc, Uint> bits;

        fenwick_set(): Base(popcount_array<maxc, Uint>{}) {}
        fenwick_set(auto &&range): fenwick_set() {
            for(auto x: range) {
                Base::data[x / word + 1] += 1;
                if(!bits.test(x)) {
                    sz++;
                    bits.flip(x);
                }
            }
            Base::to_prefix_sums();
        }
        void insert(size_t x) {
            if(bits.test(x)) return;
            Base::add(x / word, 1);
            bits.flip(x);
            sz++;
        }
        void erase(size_t x) {
            if(!bits.test(x)) return;
            Base::add(x / word, -1);
            bits.flip(x);
            sz--;
        }
        size_t order_of_key(size_t x) const {
            return Base::prefix_sum(x / word) + order_of_bit(bits.word(x / word), x % word);
        }
        size_t find_by_order(size_t order) const {
            if(order >= sz) {
                return -1;
            }
            auto [x, remainder] = Base::prefix_lower_bound(order + 1);
            return x * word + kth_set_bit(bits.word(x), remainder - 1);
        }
        size_t lower_bound(size_t x) const {
            if(bits.test(x)) {return x;}
            auto order = order_of_key(x);
            return order < sz ? find_by_order(order) : -1;
        }
        size_t pre_upper_bound(size_t x) const {
            if(bits.test(x)) {return x;}
            auto order = order_of_key(x);
            return order ? find_by_order(order - 1) : -1;
        }
    };
}

#line 1 "cp-algo/util/compress_coords.hpp"


#line 1 "cp-algo/util/sort.hpp"


#line 4 "cp-algo/util/sort.hpp"
#include <algorithm>
#include <numeric>
#include <ranges>
#line 8 "cp-algo/util/sort.hpp"
namespace cp_algo {
    template<size_t maxc>
    void count_sort(auto &a, auto &&proj = std::identity{}) {
        std::array<int, maxc> cnt = {};
        for(auto &x: a) {
            cnt[proj(x)]++;
        }
        std::partial_sum(begin(cnt), end(cnt), begin(cnt));
        auto res = a;
        for(auto const& it: a | std::views::reverse) {
            res[--cnt[proj(it)]] = it;
        }
        a = std::move(res);
    }

    void radix_sort(auto &a) {
        if(empty(a)) {
            return;
        }
        auto mx = std::ranges::max(a);
        with_bit_floor(size(a), [&]<size_t floor>() {
            constexpr int base = std::min<size_t>(floor, 1 << 16);
            for(int64_t i = 1; i <= mx; i *= base) {
                count_sort<base>(a, [&](auto x) {
                    return x / i % base;
                });
            }
        });
    }
}

#line 5 "cp-algo/util/compress_coords.hpp"
namespace cp_algo {
    // coords is a range of reference_wrapper<T>
    auto compress_coords(auto &coords) {
        std::vector<int> original;
        original.reserve(size(coords));
        radix_sort(coords);
        int idx = -1, prev = -1;
        for(auto &x: coords) {
            if(x != prev) {
                idx++;
                prev = x;
                original.push_back(x);
            }
            x.get() = idx;
        }
        return original;
    }
}

#line 6 "verify/structures/fenwick/ordered_set.test.cpp"
#include <bits/stdc++.h>

using namespace std;
using cp_algo::structures::fenwick_set;

void solve() {
    int n, q;
    cin >> n >> q;
    vector a(n, 0);
    vector<reference_wrapper<int>> coords;
    for(auto &it: a) {
        cin >> it;
        coords.push_back(ref(it));
    }
    vector queries(q, pair{0, 0});
    for(auto &[t, x]: queries) {
        cin >> t >> x;
        if(t != 2) {
            coords.push_back(ref(x));
        }
    }
    auto values = cp_algo::compress_coords(coords);
    const int maxc = 1e6;
    fenwick_set<maxc> me(a);
    for(auto [t, x]: queries) {
        if(t == 0) {
            me.insert(x);
        } else if(t == 1) {
            me.erase(x);
        } else if(t == 2) {
            int res = me.find_by_order(x-1);
            cout << (res == -1 ? -1 : values[res]) << '\n';
        } else if(t == 3) {
            cout << me.order_of_key(x+1) << '\n';
        } else if(t == 4) {
            int res = me.pre_upper_bound(x);
            cout << (res == -1 ? -1 : values[res]) << '\n';
        } else if(t == 5) {
            int res = me.lower_bound(x);
            cout << (res == -1 ? -1 : values[res]) << '\n';
        }
    }
}

signed main() {
    //freopen("input.txt", "r", stdin);
    ios::sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    while(t--) {
        solve();
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	5 ms	4 MB
g++	example_01	AC	5 ms	4 MB
g++	max_random_00	AC	90 ms	15 MB
g++	max_random_01	AC	89 ms	15 MB
g++	max_random_02	AC	70 ms	16 MB
g++	max_random_03	AC	65 ms	16 MB
g++	max_random_04	AC	82 ms	8 MB
g++	max_random_05	AC	89 ms	16 MB
g++	max_random_06	AC	100 ms	16 MB
g++	max_random_07	AC	98 ms	16 MB
g++	max_random_08	AC	98 ms	15 MB
g++	max_random_09	AC	96 ms	15 MB
g++	max_random_10	AC	71 ms	16 MB
g++	max_random_11	AC	67 ms	16 MB
g++	max_random_12	AC	84 ms	8 MB
g++	max_random_13	AC	93 ms	16 MB
g++	max_random_14	AC	113 ms	16 MB
g++	max_random_15	AC	106 ms	16 MB
g++	max_random_16	AC	137 ms	25 MB
g++	max_random_17	AC	143 ms	25 MB
g++	max_random_18	AC	119 ms	26 MB
g++	max_random_19	AC	113 ms	26 MB
g++	max_random_20	AC	148 ms	18 MB
g++	max_random_21	AC	136 ms	26 MB
g++	max_random_22	AC	151 ms	26 MB
g++	max_random_23	AC	147 ms	26 MB
g++	max_random_24	AC	150 ms	25 MB
g++	max_random_25	AC	159 ms	25 MB
g++	max_random_26	AC	122 ms	26 MB
g++	max_random_27	AC	119 ms	26 MB
g++	max_random_28	AC	146 ms	18 MB
g++	max_random_29	AC	139 ms	26 MB
g++	max_random_30	AC	162 ms	28 MB
g++	max_random_31	AC	168 ms	26 MB
g++	small_00	AC	65 ms	14 MB
g++	small_01	AC	72 ms	15 MB
g++	small_02	AC	71 ms	15 MB

Back to top page