Ghost Town problem solved with Treap Data Structure

In this article at OpenGenus, we will discuss the Ghost town problem. This involves the concept of Treap data structure.

Pre-requisites

• Maths
• Treaps
• Dynamic memory allocation

Problem Statement

Here you are given n numbers which has to be stored in a multiset. Then there will be q number of queries. There are 3 types of queries:

• 1 x: If a is the count of the elements smaller than or equal to x. Add x + a into the multiset.
• 2 y: Give the number of numbers in the multiset which are less than or equal to y.
• 3 z: Give the zth smallest number in the multiset. If numbers are repeating in the multiset then each of them has to be counted as an individual number.

Input
The first line has n and q, the number of numbers and the number of queries
Next line will have n numbers, and the rest q lines will be of the format:
Type x, where Type will be of a number 1, 2, 3 and x will be a natural number.

Output
Print the output integer for every query type 2 and 3.

Example:
Input:
10 20
7 35 44 25 15 10 21 42 12 33
1 6
1 39
2 47
2 96
1 29
2 40
3 27
3 5
1 22
1 44
3 32
1 28
3 2
2 39
3 23
2 31
1 13
1 50
3 38
2 26

Output:
11
12
10
invalid
15
invalid
7
12
invalid
8
invalid
8

Observation

The following observations can be drawn from the problem statement:

• Here the problem deals with maintaining a multiset of n numbers and handling q queries.
• The multiset stores duplicates and those numbers have to be considered.
• There are three types of operation we will be doing: inserting elements into the multiset, counting numbers in the multiset that are smaller than or equal to a given number, and the last one is to finding the zth smallest number in the multiset.
• For type 1 query, given a number x, the program needs to add x and the count of elements smaller than or equal to x to the multiset.
• For type 2 query, given a number y, the program needs to count the number of elements in the multiset that are smaller than or equal to y.
• For type 3 query, given a number z, the program needs to find the zth smallest number in the multiset. If there are duplicates, they are counted as separate elements.
• If z exceeds the number of elements in the multiset, the output for this query should be "invalid".

Brute force approach

The brute force approach for the given problem would be to maintain an array of integers and perform the following operations:

• Insert element x at the end of the array.
• Delete element x from the array.
• Find the kth smallest element in the array.

For the first operation, we can simply add the element to the end of the array.

For the second operation, we can search for the index of the element to be deleted in the array, and then shift all the elements to the right of that index one position to the left.

For the third operation, we can first sort the array in non-decreasing order and then return the kth element of the sorted array.

The time complexity of this approach would be O(qn log n) where q is the number of queries and n is the maximum size of the array, since we need to sort the array for each kth element query, which takes O(n log n) time.

Now lets discuss the optimized approach using treaps.

Approach

Intuative steps
Here we will be making use of treap data structure to solve the problem.

• In the treap, each node represents an element of the array, and is assigned a random priority value during the creation of the node.
• The priorities of the nodes follows the heap property, i.e, the priority of each node is greater than the priority of its children.
• The values stored in the nodes follow the binary search tree property, i.e., the left child of a node contains values less than or equal to the value of the node, and the right child contains values greater than or equal to the value of the node.
• The treap is maintained such that the array is represented by the values stored in the nodes, and the position of the value in the array is represented by the size of the left subtree of the node. The size of a node is the number of nodes in its subtree.
• The insert operation in treap is performed by splitting the treap into two parts based on the value to be inserted, creating a new node with the value to be inserted and a random priority, and then merging the three parts back together.
• The delete operation is performed by splitting the treap into two parts based on the value to be deleted, finding the size of the left subtree of the node to be deleted, and then merging the two parts back together.
• The find kth operation is performed by recursively traversing the treap and returning the value stored in the node with the kth size.

Explained
In the intuitive approach we have discussed the outline of what has to be done. Now lets look it in detail.

• In the starting we declare the structure of the node in the treap. The node contains a priority, a size, a value, and two pointers to the left and right nodes. The priority is a random number assigned to each node when it is created. The size of the node represents the number of nodes in the subtree rooted at that node. The value is the value of the element stored in that node.
• After that we will define two helper functions sz() and upd_sz(). sz() returns the size of the tree rooted at a given node, or 0 if the node is null. upd_sz() updates the size of the node based on the sizes of its left and right subtrees.
• Now we will make the most important functions of the program which are split() and merge().
• split() takes a treap and a key, and splits the treap into two subtrees, l and r, such that all the elements in l are less than or equal to the key, and all the elements in r are greater than the key. The function works by recursively calling itself on the left or right subtree until it reaches a null node, and then merging the resulting subtrees. The upd_sz() is called after each split to update the size of the nodes.
• merge() takes two treaps, l and r, and merges them into a single treap t. The function works by recursively calling itself on the right or left subtree until it reaches a null node, and then merging the resulting subtrees. The merge operation takes advantage of the priorities assigned to the nodes to maintain the heap property of the treap, which ensures that the maximum priority element is at the root of the treap. The upd_sz() is called after each merge to update the size of the nodes.
• After initializing is done we will take in the values of n, q and then insert these values using the insert() then we will enter a loop which does the following:
  • If the query type is 1, it performs an insert operation. First, it reads an integer x from the input. It then splits the root of the treap into two parts, where the left part contains all the elements less than or equal to x, and the right part contains all the elements greater than x. It then calculates the number of elements in the left part and adds x to it to get the number of elements less than or equal to x. It then merges the two parts of the treap back together and inserts a new node with the number of elements less than or equal to x.
  • If the query type is 2, it performs a delete operation. First, it reads an integer x from the input. It then splits the root of the treap into two parts, where the left part contains all the elements less than x, and the right part contains all the elements greater than or equal to x. It then calculates the size of the left part and prints it as the answer. It then merges the two parts of the treap back together.
  • If the query type is 3, it performs a find operation. First, it reads an integer x from the input. If the size of the treap is less than x, it prints "invalid". Otherwise, it finds the xth smallest element in the treap using the find_kth() function, and prints it to the output.
• When all queries are over we exit the loop and the program ends.

Code

#include<bits/stdc++.h>

using namespace std;


struct node
{
    int prior,size;
    int val; //value stored in the array

    int ans;
    struct node *l,*r;
};

typedef node* pnode;

int sz(pnode t)
{
    return t?t->size:0;
}

void upd_sz(pnode t)
{
    if(t)t->size=sz(t->l)+1+sz(t->r);
}

void split(pnode t,pnode &l,pnode &r,int key){
    if(!t)l=r=NULL;
    else if(t->val<=key)split(t->r,t->r,r,key),l=t;
    else split(t->l,l,t->l,key),r=t;
    upd_sz(t);
}

void merge(pnode &t,pnode l,pnode r){
    if(!l || !r)t=l?l:r;
    else if(l->prior > r->prior)merge(l->r,l->r,r),t=l;
    else merge(r->l,l,r->l),t=r;
    upd_sz(t);
}

 pnode init(int val){
    pnode ret = (pnode)malloc(sizeof(node));
    ret->val=val;ret->size=1;ret->prior=rand();ret->l=ret->r=NULL;
    return ret;
}

// kth value finding in treap
int find_kth(pnode t, int val)
{
    if(sz(t->l)+1==val) return t->val;
    if(sz(t->l)>=val) return find_kth(t->l,val);
    else return find_kth(t->r,val-sz(t->l)-1);
}

void insert(pnode &t,pnode it){
    if(!t) t=it;
    else if(it->prior>t->prior)split(t,it->l,it->r,it->val),t=it;
    else insert(t->val<=it->val?t->r:t->l,it);
    upd_sz(t);
}


int main()
{
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    // it will be the main root
    pnode root=NULL;

        int q,c,x,y,n;
        cin>>n>>q;
        while(n--)
        {
            cin>>y;
            insert(root,init(y));
        }
        while(q--)
        {
        int ch;
        cin>>ch;
        // insert after finding smaller or equal elements
        if(ch==1)
        {
            cin>>x;
            pnode l,r;
            split(root,l,r,x);
            //nw->number of elements equal or smaller than x
            int nw=sz(l);
            nw+=x;
            merge(root,l,r);
            insert(root,init(nw));
        }
       // delete element (x)
        else if(ch==2)
        {
            cin>>x;
            pnode l,r;
            split(root,l,r,x);
            int ans=sz(l);
             merge(root,l,r);
            cout<<ans<<endl;
        }
        // kth element
        else
        {
            cin>>x;
            if(sz(root)<x)
                cout<<"invalid"<<endl;
                else
            cout<<find_kth(root,x)<<endl;
        }

        }
    }

Output
If the following is the input:

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

The following will be the output:

4
2
invalid

Now lets look at how this output came.

• First query was 3, 4, here the program prints the 4th smallest element which is 4 in this case.
• Second query was 1, 2, here the program will first check how many elements are smaller than 2 in the multiset. There are 2 numbers which are smaller than 2. So 4 (2 + 2) will be inserted into the multiset, which is the sum of the given number 2 and the count of the numbers less than or equal to 2.
• Third query was 2, 2 which displays 2, because there are 2 numbers less than or equal to 2.
• Fourth query was 3, 8, here the program prints the 8th smallest element, but since there are 6 elements in the multiset it will print "invalid".

Complexities

Now lets look at the complexities. The time complexity of inserting an element is O(log n) because it involves finding the position of the element in the treap by splitting the treap into two parts based on the value of the new element. After that, a new node is inserted at the appropriate position based on the priority of the node.

The time complexity of deleting an element is also O(log n) because it involves finding the position of the element in the treap and then merging the two parts of the treap after deleting the node.

The time complexity of finding the k-th smallest element in the treap is also O(log n) because it involves finding the node in the treap which has k elements to its left.

Therefore, the overall time complexity of the given code is O(q log n), where q is the number of queries and n is the number of elements in the treap.

The space complexity of the code is O(n), where n is the number of elements in the treap, as each node of the treap takes up space proportional to the number of elements in its subtree.