Сумма элементов на отрезке с Policy-Based Data Structures (ordered_set) в GNU C++

Revision ru1, by dmkz, 2022-12-29 01:09:57

Всем привет!

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 как размер текущего узла, а мы хотим использовать значение ключа).

Небольшой тест на корректность

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

Tags c++, pbds, ordered_set, сумма подотрезка, множества

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
ru1 Russian dmkz 2022-12-29 01:09:57 6292 Первая редакция перевода на Русский
en1 English dmkz 2022-12-29 00:56:16 6103 Initial revision (published)