0% found this document useful (0 votes)
11 views6 pages

Segment Tree Multiset

Multiset
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
0% found this document useful (0 votes)
11 views6 pages

Segment Tree Multiset

Multiset
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
You are on page 1/ 6

#include <bits/stdc++.

h>

using namespace std;

/**
* Segment tree node struct.
*/
struct node {
int size;
node* childL, * childR;

/**
* Constructs an empty node.
*/
node() {
size = 0;
childL = childR = this;
}

/**
* Constructs a new node.
*
* @param s the size of the node.
* @param l the left child of this node.
* @param r the right child of this node.
*/
node(int s, node* l, node* r) {
size = s;
childL = l;
childR = r;
}

/**
* Update the size of this node.
*/
void update() {
size = childL->size + childR->size;
}
};

/**
* Multiset that store integers in the range of [-N, N].
* The multiset is implemented using segment tree.
*
* Note that the multiset is 0-indexed.
*
* The most complex function in this class is done in time complexity of O(log(N)).
*/
class segment_multiset {
const int N;
node* nil, * root;

public:

/**
* Constructs a new multiset.
*
* @param N the maximum absolute value this multiset can store.
*/
segment_multiset(int N) : N(N) {
root = nil = new node();
}

/**
* Destruct this multiset.
*/
~segment_multiset() {
destroy(root);
delete nil;
}

/**
* Clears and removes all the elements in this multiset.
*/
void clear() {
destroy(root);
root = nil;
}

/**
* Returns the total number of integers in the multiset.
*
* @return the size of the multiset.
*/
int size() {
return root->size;
}

/**
* Counts the number of occurrences of the an integer in the multiset.
*
* @param val the integer to count its occurrences.
*
* @return the number of occurrences of the given integer.
*/
int count(int val) {
node* cur = root;
int l = -N, r = N;

while (l < r) {
int mid = l + (r - l) / 2;

if (val <= mid) {


cur = cur->childL;
r = mid;
} else {
cur = cur->childR;
l = mid + 1;
}
}

return cur->size;
}

/**
* Insert a number of occurrences of an integer to the multiset.
*
* @param val the integer to insert.
* @param cnt the number of occurrences to insert (note: cnt > 0).
*/
void insert(int val, int cnt = 1) {
assert(cnt > 0);
insert(root, val, cnt, -N, N);
}

/**
* Erases a number of occurrences of an integer from the multiset if exists.
*
* @param val the integer to erase.
* @param cnt the number of occurrences to erase (note: cnt > 0).
*
* @return the number of deleted elements.
*/
int erase(int val, int cnt = 1) {
assert(cnt > 0);
return erase(root, val, cnt, -N, N);
}

/**
* Returns an integer in the multiset by its index.
*
* Note that the multiset is kept sorted in non-decreasing order,
* and is 0-indexed.
*
* @param idx the index.
*
* @return the "idx" smallest element currently in the multiset.
*/
int operator[](int idx) {
if (idx < 0 || idx >= root->size) {
throw out_of_range("ERROR :: trying to access an out of range
element");
}

node* cur = root;


int l = -N, r = N;

while (l < r) {
int mid = l + (r - l) / 2;

if (idx < cur->childL->size) {


cur = cur->childL;
r = mid;
} else {
idx -= cur->childL->size;
cur = cur->childR;
l = mid + 1;
}
}

return r;
}

/**
* Returns the index of the first integer with value greater
* than or equals to the given value.
*
* Note that the multiset is kept sorted in non-decreasing order,
* and is 0-indexed.
*
* @param val the value to returns its lower bound index.
*
* @return the specified index; or {@code segment_multiset#size() + 1}
* if such integer does not exist.
*/
int lower_bound(int val) {
int ret = 0;

node* cur = root;


int l = -N, r = N;

while (l < val) {


int mid = l + (r - l) / 2;

if (val <= mid) {


cur = cur->childL;
r = mid;
} else {
ret += cur->childL->size;
cur = cur->childR;
l = mid + 1;
}
}

return ret;
}

/**
* Returns the index of the first integer with value greater
* than the given value.
*
* Note that the multiset is kept sorted in non-decreasing order,
* and is 0-indexed.
*
* @param val the value to returns its upper bound index.
*
* @return the specified index; or {@code segment_multiset#size() + 1}
* if such integer does not exist.
*/
int upper_bound(int val) {
return lower_bound(val + 1);
}

private:

/**
* Insert a number of occurrences of an integer to the multiset.
*
* @param root the current root of the segment tree.
* @param val the integer to insert.
* @param cnt the number of occurrences to insert (note: cnt > 0).
* @param l the left index of the range of the current segment.
* @param r the right index of the range of the current segment.
*/
void insert(node*& root, int val, int cnt, int l, int r) {
if (val < l || val > r) {
return;
}

if (root == nil) {
root = new node(0, nil, nil);
}

root->size += cnt;

if (l == r) {
return;
}

int mid = l + (r - l) / 2;

insert(root->childL, val, cnt, l, mid);


insert(root->childR, val, cnt, mid + 1, r);
}

/**
* Erases a number of occurrences of an integer from the multiset if exists.
*
* @param root the current root of the segment tree.
* @param val the integer to erase.
* @param cnt the number of occurrences to erase (note: cnt > 0).
* @param l the left index of the range of the current segment.
* @param r the right index of the range of the current segment.
*
* @return the number of deleted elements.
*/
int erase(node*& root, int val, int cnt, int l, int r) {
if (val < l || val > r) {
return 0;
}

if (root == nil) {
return 0;
}

if (l == r) {
return remove(root, cnt);
}

int mid = l + (r - l) / 2;

int ret = 0;

ret += erase(root->childL, val, cnt, l, mid);


ret += erase(root->childR, val, cnt, mid + 1, r);

return remove(root, ret);


}

/**
* Removes a number of occurrences of an integer from a node and destroys it
* if it becomes empty.
*
* @param root the node to remove from.
* @param cnt the number of occurrences to erase (note: cnt > 0).
*
* @return the number of deleted elements.
*/
int remove(node*& root, int cnt) {
int ret = min(cnt, root->size);

root->size -= cnt;

if (root->size <= 0) {
destroy(root);
root = nil;
}

return ret;
}

/**
* Clears the given segment tree and releases the allocated memory.
*
* @param root the root of the sub-tree to destroy.
*/
void destroy(node* root) {
if (root == nil) return;
destroy(root->childL);
destroy(root->childR);
delete root;
}
};

You might also like