Всем привет!

PBDS настолько гибкая, что мы можем использовать свои структуры node_update с __gnu_pbds::tree<...> (tree из набора policy-based data structures). В качестве примера, можно использовать interval_node_update_policy, чтобы построить Дерево Интервалов как у Кормена. Это описано здесь.

Как вы, возможно, уже знаете: tree_order_statistics_node_update, которую все используют по дефолту, хранит размер поддерева в каждой ноде. Вместо этого нам никто не мешает хранить сумму в поддереве для вычисления суммы элементов множества на отрезке за $$$O(\log{n})$$$. Итак, эта сумма на отрезке [L,R] равна set.order_of_key(R+1)-set.order_of_key(L). Окончательно, мы можем выполнять следующие операции эффективно:

1. найти сумму элементов множества на отрезке от L до R;
2. другие операции вроде insert, erase, lower_bound, upper_bound, iterators...

Моя реализация кастомной структуры node_update:

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

using namespace std;
using namespace __gnu_pbds;

template <class Node_CItr, class Node_Itr, class Cmp_Fn, class _Alloc>
struct sum_of_keys : public tree_order_statistics_node_update<Node_CItr,Node_Itr,Cmp_Fn,_Alloc>
{
    typedef int64_t metadata_type;
    // update metadata from left, right and itself:
    void operator()(Node_Itr it, Node_CItr end_it)
    {
        metadata_type x = it.m_p_nd->m_value;
        if (it.get_l_child() != end_it) // add a sum in left child:
            x += it.get_l_child().get_metadata();
        if (it.get_r_child() != end_it) // add a sum in right child:
            x += it.get_r_child().get_metadata();
        const_cast<metadata_type&>(it.get_metadata()) = x;
    }
    // find sum of keys which is less than given `key`:
    int64_t order_of_key(const auto &key) const
    {
        const auto end_it = node_end();
        int64_t ord = 0;
        for (auto it = node_begin(); it != end_it; )
        {
            auto currKey = it.m_p_nd->m_value;
            auto left = it.get_l_child();
            if (Cmp_Fn()(key, currKey))
                it = left; // go to left child
            else if (Cmp_Fn()(currKey, key))
            {
                ord += currKey + ((left == end_it) ? 0 : left.get_metadata());
                it = it.get_r_child(); // go to right child
            }
            else // stay here
                return ord += (left == end_it) ? 0 : left.get_metadata();
        }
        return ord;
    }
    virtual Node_CItr node_begin() const = 0;
    virtual Node_CItr node_end() const = 0;
};

template<typename T, typename Comp = std::less<T>>
using OrderedSetSum = tree<T, null_type, Comp, rb_tree_tag, sum_of_keys>;

Вам нужно определить тип который вы хотите хранить в каждом узле дерева как metadata_type. После этого, вам нужно реализовать свой собственный operator() для того, чтобы смёрджить значения из левого и правого поддерева. К сожалению, мы не можем использовать уже реализованный order_of_key из наследуемого класса tree_order_statistics_node_update, потому что автор захардкодил значения для текущего узла (он использует 1 как размер текущего узла, а мы хотим использовать значение ключа).

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

using namespace std;
using namespace __gnu_pbds;

template <class Node_CItr, class Node_Itr, class Cmp_Fn, class _Alloc>
struct sum_of_keys : public tree_order_statistics_node_update<Node_CItr,Node_Itr,Cmp_Fn,_Alloc>
{
    typedef int64_t metadata_type;
    // update metadata from left, right and itself:
    void operator()(Node_Itr it, Node_CItr end_it)
    {
        metadata_type x = it.m_p_nd->m_value;
        if (it.get_l_child() != end_it)
            x += it.get_l_child().get_metadata();
        if (it.get_r_child() != end_it)
            x += it.get_r_child().get_metadata();
        const_cast<metadata_type&>(it.get_metadata()) = x;
    }
    // find sum of keys which is less than given `key`:
    int64_t order_of_key(const auto &key) const
    {
        const auto end_it = node_end();
        int64_t ord = 0;
        for (auto it = node_begin(); it != end_it; )
        {
            auto currKey = it.m_p_nd->m_value;
            auto left = it.get_l_child();
            if (Cmp_Fn()(key, currKey))
                it = left; // go to left child
            else if (Cmp_Fn()(currKey, key))
            {
                ord += currKey + ((left == end_it) ? 0 : left.get_metadata());
                it = it.get_r_child(); // go to right child
            }
            else // return
                return ord += (left == end_it) ? 0 : left.get_metadata();
        }
        return ord;
    }
    virtual Node_CItr node_begin() const = 0;
    virtual Node_CItr node_end() const = 0;
};

template<typename T, typename Comp = std::less<T>>
using OrderedSetSum = tree<T, null_type, Comp, rb_tree_tag, sum_of_keys>;

int main()
{
    OrderedSetSum<int> set;
    for (int i = 0; i < 10; i++) {
        set.insert(i * 10);
    }
    // set: {0, 10, 20, 30, 40, 50, 60, 70, 80, 90}
    auto slow = [&](int L, int R) {
        int64_t answ{};
        for (auto it : set)
            if (L <= it && it <= R)
                answ += it;
        return answ;
    };
    auto fast = [&](int L, int R) {
        return set.order_of_key(R+1) - set.order_of_key(L);
    };
    for (int i = 0; i < 100; i++) {
        for (int j = i; j < 100; j++) {
            assert(slow(i,j) == fast(i,j));
        }
    }
    std::cout << "All tests have been passed!" << std::endl;
    return 0;
}

Буду продолжать исследовать что может быть сделано с tree из Policy-Based Data Structures и держать вас в курсе. Пожалуйста, делитесь своими идеями в комментариях или присылайте ссылки если это уже сделано кем-то. Мой предыдущий блог, где я объединил lower_bound с order_of_key здесь.