I have merged lower_bound and order_of_key into a single method in C++ ordered_set

#	User	Rating
1	tourist	3985
2	jiangly	3814
3	jqdai0815	3682
4	Benq	3529
5	orzdevinwang	3526
6	ksun48	3517
7	Radewoosh	3410
8	hos.lyric	3399
9	ecnerwala	3392
9	Um_nik	3392

#	User	Contrib.
1	cry	169
2	maomao90	162
2	Um_nik	162
4	atcoder_official	161
5	djm03178	158
6	-is-this-fft-	157
7	adamant	155
8	awoo	154
8	Dominater069	154
10	luogu_official	150

Hi everyone!

I took an ordered_set in C++ (i.e. __gnu_pbds::tree<...> from GNU C++ Standard Template Library with implemented Policy-Based Data Structures) and added 6 new methods to derived class:

lower_bound_with_order -> iterator to minimal element such >= than given key + its index in the set (its order);
upper_bound_with_order -> iterator to minimal element such > than given key + its index in the set (its order);
count_less -> number of elements which are less than given element;
count_less_equal -> number of elements which are less or equal than given element;
count_greater -> number of elements which are greater than given element;
count_equal -> just a number of elements which are equal to given element.

Advantages of (1) and (2): it is implemented in a single traversal of a tree ($$$\log{N}$$$). If you will call lower_bound firstly, and call order_of_key secondly, then you will spend in two times more time for doing it ($$$\log{N} + \log{N}$$$) instead.

Source code

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

template<typename T, typename Comp = std::less<T>>
using OrdTreeBase = __gnu_pbds::tree<T, __gnu_pbds::null_type, Comp,
                                     __gnu_pbds::rb_tree_tag,
                                     __gnu_pbds::tree_order_statistics_node_update>;


template<typename T, typename Comp = std::less<T>>
struct OrderedSet : public OrdTreeBase<T, Comp>
{
    using Base = OrdTreeBase<T, Comp>;
    
    OrderedSet(const auto &...pack)
        : Base(pack...) { }
    
    size_t count_less(const auto &key) const {
        return this->order_of_key(key);
    }
    
    size_t count_less_equal(const auto &key) const {
        return upper_bound_with_order(key).second;
    }
    
    size_t count_greater(const auto &key) const {
        return this->size() - count_less_equal(key);
    }
    
    size_t count_equal(const auto &key) const {
        return count_less_equal(key) - count_less(key);
    }
    
    template<typename Compare = Comp>
    std::pair<typename Base::const_iterator, size_t>
    lower_bound_with_order(const auto &key) const
    {
        auto prev = this->m_p_head;
        auto curr = this->m_p_head->m_p_parent;
        size_t order{};
        while (curr != 0)
            if (Compare()(curr->m_value, key)) {
                order += curr->m_metadata;
                if ((curr = curr->m_p_right))
                    order -= curr->m_metadata;
            } else {
                prev = curr;
                curr = curr->m_p_left;
            }
        return std::make_pair(prev, order);
    }
    
    template<typename Compare = Comp>
    std::pair<typename Base::const_iterator, size_t>
    upper_bound_with_order(const auto &key) const
    {
        auto prev = this->m_p_head;
        auto curr = this->m_p_head->m_p_parent;
        size_t order{};
        while (curr != 0)
            if (Compare()(key, curr->m_value)) {
                prev = curr;
                curr = curr->m_p_left;
            } else {
                order += curr->m_metadata;
                if ((curr = curr->m_p_right))
                    order -= curr->m_metadata;
            }
        return std::make_pair(prev, order);
    }
    
};

const int INF = (int)1e9+7;

template<typename T, typename Comp>
void test(const char * name) {
    printf("--------------------------------------\n");
    printf("test: %s\n", name);
    std::vector<T> vec;
    for (int i = 0; i < 10; i += 2)
        vec.push_back(i);
    OrderedSet<T, Comp> set(vec.begin(), vec.end());
    printf("set:");
    for (auto it : set)
        printf(" %d", it);
    printf("\n");
    printf("lower_bound_with_order:\n");
    for (int i = -1; i < 10; i++) {
        auto [it, order] = set.lower_bound_with_order(i);
        if (it != set.end())
            assert(set.order_of_key(*it) == order);
        else
            assert(order == set.size());
        int value = (it == set.end()) ? INF : *it;
        printf("\t%d: lb = %d, order = %lu\n", i, value, order);
    }
    printf("upper_bound_with_order:\n");
    for (int i = -1; i < 10; i++) {
        auto [it, order] = set.upper_bound_with_order(i);
        if (it != set.end())
            assert(set.order_of_key(*it) == order);
        else
            assert(order == set.size());
        int value = (it == set.end()) ? INF : *it;
        printf("\t%d: up = %d, order = %lu\n", i, value, order);
    }
}

int main() {
    test<int, std::less<int>>("test sorted");
    test<int, std::greater<int>>("test reversed");
}

Result of testing

test: test sorted
set: 0 2 4 6 8
lower_bound_with_order:
        -1: lb = 0, order = 0
        0: lb = 0, order = 0
        1: lb = 2, order = 1
        2: lb = 2, order = 1
        3: lb = 4, order = 2
        4: lb = 4, order = 2
        5: lb = 6, order = 3
        6: lb = 6, order = 3
        7: lb = 8, order = 4
        8: lb = 8, order = 4
        9: lb = 1000000007, order = 5
upper_bound_with_order:
        -1: up = 0, order = 0
        0: up = 2, order = 1
        1: up = 2, order = 1
        2: up = 4, order = 2
        3: up = 4, order = 2
        4: up = 6, order = 3
        5: up = 6, order = 3
        6: up = 8, order = 4
        7: up = 8, order = 4
        8: up = 1000000007, order = 5
        9: up = 1000000007, order = 5

test: test reversed
set: 8 6 4 2 0
lower_bound_with_order:
        -1: lb = 1000000007, order = 5
        0: lb = 0, order = 4
        1: lb = 0, order = 4
        2: lb = 2, order = 3
        3: lb = 2, order = 3
        4: lb = 4, order = 2
        5: lb = 4, order = 2
        6: lb = 6, order = 1
        7: lb = 6, order = 1
        8: lb = 8, order = 0
        9: lb = 8, order = 0
upper_bound_with_order:
        -1: up = 1000000007, order = 5
        0: up = 1000000007, order = 5
        1: up = 0, order = 4
        2: up = 0, order = 4
        3: up = 2, order = 3
        4: up = 2, order = 3
        5: up = 4, order = 2
        6: up = 4, order = 2
        7: up = 6, order = 1
        8: up = 6, order = 1
        9: up = 8, order = 0

You can take it to your CP template, if you want. I will try to do same for a ordered_multiset. I'm trying ordered_set with a comparator std::less_equal instead of std::less...

dmkz's blog