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 $$$O(\log{n})$$$. 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:
- sum of subset of set with items from
L
toR
; - 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.