Kth Smallest Number in the Array

Innoskrit August 18, 2022 Innoskrit Masterclass DSA Sheet, Priority Queue (Heap)

Problem Statement

Given an unsorted array of numbers, find Kth smallest number in it.
Please note that it is the Kth smallest number in the sorted order, not the Kth distinct element.

Example:

Input 1: [1, 5, 12, 2, 11, 5], K = 3
Output 1: 5
Explanation 1: The 3rd smallest number is '5', as the first two smaller numbers are [1, 2].

Input 2: [1, 5, 12, 2, 11, 5], K = 4
Output 2: 5
Explanation 2: The 4th smallest number is '5', as the first three small numbers are [1, 2, 5].

Code

Java

C++

Python

import java.util.*;

class Innoskrit {

  public static int findKthSmallestNumber(int[] nums, int k) {
    PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>((n1, n2) -> n2 - n1);
    for (int i = 0; i < k; i++)
      maxHeap.add(nums[i]);

    for (int i = k; i < nums.length; i++) {
      if (nums[i] < maxHeap.peek()) {
        maxHeap.poll();
        maxHeap.add(nums[i]);
      }
    }

    return maxHeap.peek();
  }

  public static void main(String[] args) {
    int result = Innoskrit.findKthSmallestNumber(new int[] { 1, 5, 12, 2, 11, 5 }, 3);
    System.out.println(result);

    result = Innoskrit.findKthSmallestNumber(new int[] { 1, 7, 12, 2, 11, 5 }, 4);
    System.out.println(result);
      
  }
  
}

using namespace std;

#include <iostream>
#include <queue>
#include <vector>

class Innoskrit {
 public:
  static int findKthSmallestNumber(const vector<int> &nums, int k) {
    priority_queue<int> maxHeap;

    for (int i = 0; i < k; i++) {
      maxHeap.push(nums[i]);
    }

    for (int i = k; i < nums.size(); i++) {
      if (nums[i] < maxHeap.top()) {
        maxHeap.pop();
        maxHeap.push(nums[i]);
      }
    }

    return maxHeap.top();
  }
};

int main(int argc, char *argv[]) {
  int result = Innoskrit::findKthSmallestNumber(vector<int>{1, 5, 12, 2, 11, 5}, 3);
  cout << result << endl;

  result = Innoskrit::findKthSmallestNumber(vector<int>{1, 7, 12, 2, 11, 5}, 4);
  cout << result << endl;
}

from heapq import *

def find_Kth_smallest_number(nums, k):
  
  maxHeap = []

  for i in range(k):
    heappush(maxHeap, -nums[i])

  for i in range(k, len(nums)):
    if -nums[i] > maxHeap[0]:
      heappop(maxHeap)
      heappush(maxHeap, -nums[i])

  return -maxHeap[0]

def main():

  print(find_Kth_smallest_number([1, 5, 12, 2, 11, 5], 3))

  print(find_Kth_smallest_number([1, 7, 12, 2, 11, 5], 4))

main()

Output:

5
7

Time and Space Complexity

Time Complexity:

The time complexity of this algorithm is O(K*logK+(N-K)*logK), which is asymptotically equal to O(N*logK).

Space Complexity: The space complexity will be O(K) since we need to store the top ‘K’ numbers in the heap.