Kth order-Segment tree problem in C++

Now we will be going to solve the segment tree question in C++.

Problem description

Given n integers, and an array an of n integers, and q queries. The queries are of three types.

Type 1 – add integer v  to the array.

Type 2 – remove the Kth smallest integer from the array.

Type 3 – Output the Kth smallest integer of the array.

Input

First line contains integers n and q. Second line contains n integers a1,a2,a3…an,(1<=ai<=n).

Next q line contains queries of the following type:-

1 v (add integer v to array)

2 k (remove kth smallest integer)

3 k (output kth smallest integer)

It is guaranteed that k<=currentsizeofarray

For query of type 3, output corresponding answer.

Constraints

0<=n<=500000

0<=q<=500000

1<=v<=n

Input

5 5
4 2 4 2 2
2 5
3 2
1 2
2 3
3 4

Output

2
4

Approach

We will maintain a segment tree. Each node of the segment tree will store the number of integers in that range. (For example – root (here seg[1] is root) will store a number of integers from range 1 to n currently present in array). For inserting an integer x, we will increment the value of the node corresponding to range x to x, and then update all ancestors associated with this node. This will take log(n) operations. For getting the K-th value, we will start from the root, and see its left and right child. Let left child value be V. if V is greater/equal to K, we will look in left child for K-th value, else we will look in a right child for (K-V)-th value. For removing the K-th value, we will perform the above steps and decrement the node values each time.

Segment tree

A Segment Tree is a data structure that allows answering range queries over an array effectively, it is used to give the answer of a particular range of input [L….R]. normally we used to calculate the answer to this range by iterating from L to R  but the segment tree modifies this approach to calculate the answer ex:- finding minimum element or sum of elements in [L….R]  we don’t need to iterate over L to R we can use segment tree to answer this query in O(log(n)) time(in the worst case we had to traverse till depth os tree which is log(n) ). it is represented in tree form with maximum height up to log(n) where n=number of nodes and each node contains the answer to the query. Along with answering such queries the Segment Tree also allows modifying the array element data by replacing one element or even change the elements of a whole subsegment. so we can say that a Segment Tree is a data structure from which a huge number of problems can be solved with it and the standard Segment Tree requires 4n4n vertices for working on an array of size n.

Implementation 

#include<bits/stdc++.h>
using namespace std;
int n,q,t,k,seg[1048576];
void ins(int x,int s=1,int e=n,int i=1){
    ++seg[i];
    if(s==e) return;
    int md=(s+e)/2;
    (x>md)?ins(x,md+1,e,2*i+1):ins(x,s,md,2*i);
}
void rem(int k,int s=1,int e=n,int i=1){
    --seg[i];
    if(s==e) return;
    int md=(s+e)/2;
    (seg[2*i]>=k)?rem(k,s,md,2*i):rem(k-seg[2*i],md+1,e,2*i+1);
}
int que(int x,int s=1,int e=n,int i=1){
    if(s==e) return s;
    int md=(s+e)/2;
    return seg[2*i]>=x?que(x,s,md,2*i):que(x-seg[2*i],md+1,e,2*i+1); 
}
int main(){
    ios::sync_with_stdio(0);cin.tie(0);
    cin>>n>>q;
    for(int i=1;i<=n;++i) cin>>k,ins(k);
    while(q--){
        cin>>t>>k;
        if(t==1) ins(k);
        else if(t==2) rem(k);
        else cout<<que(k)<<'\n';
    }
}

Sample case

Let us suppose n=5, and insert integer 2, Our array will become (1,2,2,3,5) And now let us suppose if we have to find 4th integer, For removing query, we will follow the above procedure and decrement the value of each node visited. Time complexity O(nlog(n)+qlog(n)) and space complexity O(n).

Hope you understand the implementation and algorithm used! Thank you.