K Largest Numbers in the Array

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

Problem Statement

Given an unsorted array of numbers, find the ‘K’ largest numbers in it.

Example:

Input 1: arr: [3, 1, 5, 11, 4, 10], k = 3
Output 1: [5, 11, 10]

Input 1: arr: [3, 1, -5, -11, 4, 10], k = 4
Output 1: [1, 3, 4, 10]

Note: Order of elements doesn't matter.

Code

Java

C++

Python

import java.util.*;

class Innoskrit {

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

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

    return new ArrayList<>(minHeap);
  }

  public static void main(String[] args) {
    List<Integer> result = Innoskrit.findKLargestNumbers(new int[] {3, 1, 5, 11, 4, 10}, 3);
    System.out.println(result);

    result = Innoskrit.findKLargestNumbers(new int[] {3, 1, -5, -11, 4, 10}, 4);
    System.out.println(result);
  }
}

using namespace std;
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>

class Innoskrit {
 public:
  struct greater {
    bool operator()(const int& a, const int& b) const { return a > b; }
  };
  static vector<int> findKLargestNumbers(const vector<int>& nums, int k) {
    
    vector<int> minHeap(nums.begin(), nums.begin() + k);
    make_heap(minHeap.begin(), minHeap.end(), greater());

    for (int i = k; i < nums.size(); i++) {
      if (nums[i] > minHeap.front()) {
        pop_heap(minHeap.begin(), minHeap.end(), greater());
        minHeap.pop_back();
        minHeap.push_back(nums[i]);
        push_heap(minHeap.begin(), minHeap.end(), greater());
      }
    }

    return minHeap;
  }
};

int main(int argc, char* argv[]) {
  vector<int> result = Innoskrit::findKLargestNumbers(vector<int>{3, 1, 5, 11, 4, 10}, 3);
  for (auto num : result) {
    cout << num << " ";
  }
  cout << endl;

  result = Innoskrit::findKLargestNumbers(vector<int>{3, 1, -5, -11, 4, 10}, 4);
  for (auto num : result) {
    cout << num << " ";
  }
  cout << endl;
}

from heapq import *

def find_k_largest_numbers(nums, k):
  minHeap = []
  
  for i in range(k):
    heappush(minHeap, nums[i])

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

  return minHeap
  
  
if __name__ == "__main__":
  print(find_k_largest_numbers([3, 1, 5, 11, 4, 10], 3))
  print(find_k_largest_numbers([3, 1, -5, -11, 4, 10], 4))

Output:

[5, 11, 10]
[1, 3, 4, 10]

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.