Sum of keys on segment with Policy-Based Data Structures (ordered_set) in GNU C+...

Hello again!

We can use custom node_update structure with our __gnu_pbds::tree<...> (tree from set of policy-based data structures). For example, there is interval_node_update_policy, using of which allows us to build the Interval Tree. It is described here.

As you may already know, that tree_order_statistics_node_update stores a size of subtree for each node. Instead of it, we can store sum of keys in subtree for calculating sum of keys on segment in . So, sum of keys on segment [L,R] is just set.order_of_key(R+1)-set.order_of_key(L). Finally, we can perform next set of operations efficiently:

1. sum of subset of set with items from L to R;
2. another operations like insert, erase, lower_bound, upper_bound, iterators...

Here is my implementation of sum of keys on segment:

#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>;

As you can see, you should define which data type you want to store in node of tree as metadata_type. After this, you should implement your operator() for merging values from left and right subtrees. Unfortunately, we can not use order_of_key from base class tree_order_statistics_node_update, because value of current node is hardcoded (author uses 1 as size of single node, but we should use value, which is written in node.

I will continue an investigation of what can be done with a tree from Policy-Based Data Structures. Please, write your ideas in comments or send links if it is already done by someone. My previous blog, where I have merged lower_bound with order_of_key is here.