Find K-th smallest element in an array for multiple queries

Given an array arr[] of size N and an array Q[][] consisting of M queries that needs to be processed on the given array. It is known that these queries can be of the following two types:

• Type 1: If Q = 1, then add an element in the array {type, element_to_add}.
• Type 2: If Q = 2, then print the K-th smallest element in the array. {type, k}

The task is to perform the given queries such that the following constraints are given:

• 1 ? Number of queries ? 1000000.
• 1 ? N ? 1000000.
• 1 ? arr[i] ? 100000000. And, the array can contain duplicate entries.

Examples:

Input: arr[] = {1, 2, 3, 4, 5, 6}, Q[][] = {{1, 7}, {2, 4}, {1, 5}, {2, 6}} Output: 4 5 Explanation: The first query is used to add 7 in the array. The array arr[] becomes: {1, 2, 3, 4, 5, 6, 7} The second query is to find 4^th smallest element. It is 4 in this case. The third query is used to add 5 in the array. The array arr[] becomes: {1, 2, 3, 4, 5, 5, 6, 7} The fourth query is to find 6^th smallest element. It is 5 in this case. Input: arr[] = {2, 4, 2, 1, 5, 6, 2, 4}, Q[][] = {{1, 3}, {2, 2}, {1, 10}, {2, 7}} Output: 2 4

Naive Approach: The naive approach for this problem is to add the element in the array and sort the array for every first type of query. And, whenever the second type of query occurs, print the K-th element in the sorted array. Time complexity: O(M * (Nlog(N))) where M is the number of queries and N is the size of the array. Efficient Approach: The idea is to use a policy-based data structure. For this problem, we can use order_of_key datastructure to find the K-th smallest element in the array.

• The order_of_key() is a builtin function of Ordered Set which is a Policy Based Data Structure in C++. Policy-based data structures are not part of the C++ standard library but g++ compiler supports them. This data structure finds the K-th smallest element in logarithmic time complexity.
• However, this data structure doesn’t allow duplicate keys. In order to use the data structure for duplicate elements, a pair is used.
• We create a pair of the given element and index number to insert it in the policy-based data structure.
• The pairs are first sorted by comparing the first element than the second. For example, (2, 1) is ordered before (2, 7).

Below is the implementation of the above approach: 

cpp

// C++ program to find K-th
// smallest element in an array
// for multiple queries

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

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

template <class T>

// Defining the policy-based data
// structure
using oset
    = tree<T, null_type,
        less<T>, rb_tree_tag,
        tree_order_statistics_node_update>;

// Function to find K-th
// smallest element in an array
// for multiple queries
void Operations(int arr[], int n,
                int query[][2], int k)
{

    // Declare the data structure that
    // stores the pairs
    oset<pair<int, int> > s;

    int j = 0;
    // Loop to insert the initial array
    // elements into the set
    for (j = 0; j < n; j++) {
        s.insert({ arr[j], j });
    }

    // Loop to process the queries
    for (int i = 0; i < k; i++) {
        if (query[i][0] == 1) {

            // Inserting the element if the
            // type of query is 1
            s.insert({ query[i][1], j });
            j++;
        }

        // Finding the K-th smallest element
        // if the type of the query is 2
        else if (query[i][0] == 2) {
            cout << (*s
                        .find_by_order(
                            query[i][1] - 1))
                        .first
                << endl;
        }
    }
}

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

    int k = 4;
    // Queries. The first element denotes
    // the type of the query
    int query[4][2] = { { 1, 3 },
                        { 2, 2 },
                        { 1, 10 },
                        { 2, 7 } };

    Operations(arr, n, query, k);

    return 0;
}

Java

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class Main {

    // Function to find K-th smallest element in an array
    // for multiple queries
    static void operations(int[] arr, int n,
                           int[][] queries, int k)
    {
        // Declare the data structure that stores the pairs
        List<Pair> s = new ArrayList<>();

        int j = 0;
        // Loop to insert the initial array elements into
        // the set
        for (j = 0; j < n; j++) {
            s.add(new Pair(arr[j], j));
        }
        Collections.sort(s);

        // Loop to process the queries
        for (int i = 0; i < k; i++) {
            if (queries[i][0] == 1) {
                // Inserting the element if the type of
                // query is 1
                s.add(new Pair(queries[i][1], j));
                j++;
                Collections.sort(s);
            }
            else if (queries[i][0] == 2) {
                // Finding the K-th smallest element if the
                // type of the query is 2
                System.out.println(
                    s.get(queries[i][1] - 1).first);
            }
        }
    }

    // Pair class to store pairs of elements and their
    // indices
    static class Pair implements Comparable<Pair> {
        int first, second;

        Pair(int first, int second)
        {
            this.first = first;
            this.second = second;
        }

        // Comparison based on the first element of the pair
        public int compareTo(Pair other)
        {
            return Integer.compare(this.first, other.first);
        }
    }

    // Driver code
    public static void main(String[] args)
    {
        int n = 8;
        int[] arr = { 2, 4, 2, 1, 5, 6, 2, 4 };

        int k = 4;
        // Queries. The first element denotes the type of
        // the query
        int[][] queries
            = { { 1, 3 }, { 2, 2 }, { 1, 10 }, { 2, 7 } };

        operations(arr, n, queries, k);
    }
}

Python3

# Function to find K-th
# smallest element in an array
# for multiple queries
def operations(arr, n, queries, k):
    # Declare the data structure that
    # stores the pairs
    s = []

    j = 0
    # Loop to insert the initial array
    # elements into the set
    for j in range(n):
        s.append((arr[j], j))
    s.sort()

    # Loop to process the queries
    for i in range(k):
        if queries[i][0] == 1:
            # Inserting the element if the
            # type of query is 1
            s.append((queries[i][1], j))
            j += 1
            s.sort()
        elif queries[i][0] == 2:
            # Finding the K-th smallest element
            # if the type of the query is 2
            print(s[queries[i][1] - 1][0])

# Driver code
if __name__ == "__main__":
    n = 8
    arr = [2, 4, 2, 1, 5, 6, 2, 4]

    k = 4
    # Queries. The first element denotes
    # the type of the query
    queries = [[1, 3], [2, 2], [1, 10], [2, 7]]

    operations(arr, n, queries, k)

C#

using System;
using System.Collections.Generic;

public class KthSmallestElement
{
    // Function to find K-th smallest element in an array
    // for multiple queries
    static void Operations(int[] arr, int n,
                           int[,] queries, int k)
    {
        // Declare the data structure that stores the pairs
        List<Pair> s = new List<Pair>();

        int j = 0;
        // Loop to insert the initial array elements into
        // the list
        for (j = 0; j < n; j++)
        {
            s.Add(new Pair(arr[j], j));
        }
        s.Sort();

        // Loop to process the queries
        for (int i = 0; i < k; i++)
        {
            if (queries[i, 0] == 1)
            {
                // Inserting the element if the type of
                // query is 1
                s.Add(new Pair(queries[i, 1], j));
                j++;
                s.Sort();
            }
            else if (queries[i, 0] == 2)
            {
                // Finding the K-th smallest element if the
                // type of the query is 2
                Console.WriteLine(s[queries[i, 1] - 1].first);
            }
        }
    }

    // Pair class to store pairs of elements and their
    // indices
    public class Pair : IComparable<Pair>
    {
        public int first, second;

        public Pair(int first, int second)
        {
            this.first = first;
            this.second = second;
        }

        // Comparison based on the first element of the pair
        public int CompareTo(Pair other)
        {
            return this.first.CompareTo(other.first);
        }
    }

    // Driver code
    public static void Main(string[] args)
    {
        int n = 8;
        int[] arr = { 2, 4, 2, 1, 5, 6, 2, 4 };

        int k = 4;
        // Queries. The first element denotes the type of
        // the query
        int[,] queries
            = { { 1, 3 }, { 2, 2 }, { 1, 10 }, { 2, 7 } };

        Operations(arr, n, queries, k);
    }
}
//this code is contributed by Prachi

Javascript

// Function to find K-th
// smallest element in an array
// for multiple queries
function operations(arr, n, queries, k) {
    // Declare the data structure that
    // stores the pairs
    let s = [];

    let j = 0;
    // Loop to insert the initial array
    // elements into the set
    for (j = 0; j < n; j++) {
        s.push([arr[j], j]);
    }
    s.sort((a, b) => a[0] - b[0]);

    // Loop to process the queries
    for (let i = 0; i < k; i++) {
        if (queries[i][0] === 1) {
            // Inserting the element if the
            // type of query is 1
            s.push([queries[i][1], j]);
            j++;
            s.sort((a, b) => a[0] - b[0]);
        } else if (queries[i][0] === 2) {
            // Finding the K-th smallest element
            // if the type of the query is 2
            console.log(s[queries[i][1] - 1][0]);
        }
    }
}

// Driver code
let n = 8;
let arr = [2, 4, 2, 1, 5, 6, 2, 4];

let k = 4;
// Queries. The first element denotes
// the type of the query
let queries = [[1, 3], [2, 2], [1, 10], [2, 7]];

operations(arr, n, queries, k);

Output

2
4

Time Complexity: O(M * log(N)), where M is the number of queries and N is the size of the array.

Auxiliary Space:  O(N+k)