Queries to find kth smallest element and point update : Ordered Set in C++

Given an array arr[] of size N and a set Q[][] containing M queries, the task is to execute the queries on the given array such that there can be two types of queries:

• Type 1: [i, x] – Update the element at i^th index to x.
• Type 2: [k] – Find the k^th smallest element in the array.

Examples:

Input: arr[] = {4, 3, 6, 2}, Q[][] = {{1, 2, 5}, {2, 3}, {1, 1, 7}, {2, 1}} Output: 5 2 Explanation: For the 1^st query: arr[] = {4, 5, 6, 2} For the 2^nd query: 3^rd smallest element would be 5. For the 3^rd query: arr[] = {7, 5, 6, 2} For the 4^th query: 1^st smallest element would be 2. Input: arr[] = {1, 0, 4, 2, 0}, Q[][] = {{1, 2, 1}, {2, 2}, {1, 4, 5}, {1, 3, 7}, {2, 1}, {2, 5}} Output: 1 0 7

Naive Approach: The naive approach for this problem is to update the i^th element in an array in constant time and use sorting to find the K^th smallest element. 

Below is the implementation of the approach:

C++

// C++ code for the approach

#include <bits/stdc++.h>
using namespace std;

// Driver code
int main()
{
    int n = 5, m = 6;
    int arr[] = { 1, 0, 4, 2, 0 };

    vector<vector<int> > query
        = { { 1, 2, 1 }, { 2, 2 }, { 1, 4, 5 },
            { 1, 3, 7 }, { 2, 1 }, { 2, 5 } };

    // Execute queries
    for (auto q : query) {
        // update
        if (q[0] == 1) { // Type 1 query
            int i = q[1] - 1; // 0-indexed
            int x = q[2];

            // update the ith index value with x
            arr[i] = x;
        }
        // find Kth smallest element
        else if (q[0] == 2) { // Type 2 query
            int k = q[1];

            // sort the array
            sort(arr, arr + n);

            // Print the kth smallest element
            cout << arr[k - 1] << endl;
        }
    }

    return 0;
}

Java

// Java code for the approach

import java.util.*;

public class GFG {
      // Driver's code
    public static void main(String[] args)
    {
        int n = 5, m = 6;
        int[] arr = { 1, 0, 4, 2, 0 };

        List<List<Integer> > query = new ArrayList<>();
        query.add(Arrays.asList(1, 2, 1));
        query.add(Arrays.asList(2, 2));
        query.add(Arrays.asList(1, 4, 5));
        query.add(Arrays.asList(1, 3, 7));
        query.add(Arrays.asList(2, 1));
        query.add(Arrays.asList(2, 5));

        // Execute queries
        for (List<Integer> q : query) {
            // update
            if (q.get(0) == 1) { // Type 1 query
                int i = q.get(1) - 1; // 0-indexed
                int x = q.get(2);

                // update the ith index value with x
                arr[i] = x;
            }
            // find Kth smallest element
            else if (q.get(0) == 2) { // Type 2 query
                int k = q.get(1);

                // sort the array
                Arrays.sort(arr);

                // Print the kth smallest element
                System.out.println(arr[k - 1]);
            }
        }
    }
}

Output

1
0
7

Time Complexity: O(M * (N * log(N))) where M is the number of queries and N is the size of the array. 
Auxiliary Space: O(1) as no extra space has been used.

Efficient Approach: The idea is to use a policy-based data structure similar to a set. Here, a tree based container is used to store the array in the form of a sorted tree such that all the nodes to the left are smaller than the root and all the nodes to the right are greater than the root. The following are the properties of the data structure:

• It is indexed by maintaining node invariant where each node contains a count of nodes in its subtree.
• Every time we insert a new node or delete a node, we can maintain the invariant in O(logN) time by bubbling up to the root.
• So the count of the node in its left subtree gives the index of that node in sorted order because the value of every node of the left subtree is smaller than the parent node.

Therefore, the idea is to follow the following approach for each query:

1. Type 1: For this query, we update the i^th element of the array. Therefore, we need to update the element both in the array and the data structure. In order to update the value in the tree container, the value arr[i] is found in the tree, deleted from the tree and the updated value is inserted back into the tree.
2. Type 2: In order to find the K^th smallest element, find_by_order(K – 1) is used on the tree as the data is a sorted data. This is similar to Binary Search operation on the sorted array.

Below is the implementation of the above approach: 

CPP

// C++ implementation of the above approach

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;

// Defining the policy based Data Structure
typedef tree<pair<int, int>,
             null_type,
             less<pair<int, int> >,
             rb_tree_tag,
             tree_order_statistics_node_update>
    indexed_set;

// Elements in the array are not unique,
// so a pair is used to give uniqueness
// by incrementing cnt and assigning
// with array elements to insert in mySet
int cnt = 0;

// Variable to store the data in the
// policy based Data Structure
indexed_set mySet;

// Function to insert the elements
// of the array in mySet
void insert(int n, int arr[])
{
    for (int i = 0; i < n; i++) {
        mySet.insert({ arr[i], cnt });
        cnt++;
    }
}

// Function to update the value in
// the data structure
void update(int x, int y)
{
    // Get the pointer of the element
    // in mySet which has to be updated
    auto it = mySet.lower_bound({ y, 0 });

    // Delete from mySet
    mySet.erase(it);

    // Insert the updated value in mySet
    mySet.insert({ x, cnt });
    cnt++;
}

// Function to find the K-th smallest
// element in the set
int get(int k)
{
    // Find the pointer to the kth smallest element
    auto it = mySet.find_by_order(k - 1);
    return (it->first);
}

// Function to perform the queries on the set
void operations(int arr[], int n,
                vector<vector<int> > query, int m)
{
    // To insert the element in mySet
    insert(n, arr);

    // Iterating through the queries
    for (int i = 0; i < m; i++) {

        // Checking if the query is of type 1
        // or type 2
        if (query[i][0] == 1) {

            // The array is 0-indexed
            int j = query[i][1] - 1;
            int x = query[i][2];

            // Update the element in mySet
            update(x, arr[j]);

            // Update the element in the array
            arr[j] = x;
        }
        else {
            int K = query[i][1];

            // Print Kth smallest element
            cout << get(K) << endl;
        }
    }
}

// Driver code
int main()
{
    int n = 5, m = 6, arr[] = { 1, 0, 4, 2, 0 };

    vector<vector<int> > query = { { 1, 2, 1 },
                                   { 2, 2 },
                                   { 1, 4, 5 },
                                   { 1, 3, 7 },
                                   { 2, 1 },
                                   { 2, 5 } };

    operations(arr, n, query, m);

    return 0;
}

Output:

1
0
7

Time Complexity: Since every operation takes O(Log(N)) time and there are M queries, the overall time complexity is O(M * Log(N)).