Leetcode

Kth Largest Element in an Array

May 26, 2022 by aAshish

Approach 1: Heap

  • Time:O(n\log k)
  • Space:O(k)

C++

 
class Solution {
 public:
  int findKthLargest(vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<>> minHeap;

    for (const int num : nums) {
      minHeap.push(num);
      if (minHeap.size() > k)
        minHeap.pop();
    }

    return minHeap.top();
  }
};

JAVA

 
class Solution {
  public int findKthLargest(int[] nums, int k) {
    Queue<Integer> minHeap = new PriorityQueue<>((a, b) -> a - b);

    for (final int num : nums) {
      minHeap.offer(num);
      while (minHeap.size() > k)
        minHeap.poll();
    }

    return minHeap.peek();
  }
}

Python

 
class Solution:
  def findKthLargest(self, nums: List[int], k: int) -> int:
    minHeap = []

    for num in nums:
      heapq.heappush(minHeap, num)
      if len(minHeap) > k:
        heapq.heappop(minHeap)

    return minHeap[0]

Approach 2: Quick Select

  • Time:O(n) \to O(n^2)
  • Space:O(1)

C++

 
class Solution {
 public:
  int findKthLargest(vector<int>& nums, int k) {
    return quickSelect(nums, 0, nums.size() - 1, k);
  }

 private:
  int quickSelect(vector<int>& nums, int l, int r, int k) {
    const int pivot = nums[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (nums[i] >= pivot)
        swap(nums[nextSwapped++], nums[i]);
    swap(nums[nextSwapped], nums[r]);

    const int count = nextSwapped - l + 1;  // # of nums >= pivot
    if (count == k)
      return nums[nextSwapped];
    if (count > k)
      return quickSelect(nums, l, nextSwapped - 1, k);
    return quickSelect(nums, nextSwapped + 1, r, k - count);
  }
};

JAVA

 
class Solution {
  public int findKthLargest(int[] nums, int k) {
    return quickSelect(nums, 0, nums.length - 1, k);
  }

  private int quickSelect(int[] nums, int l, int r, int k) {
    final int pivot = nums[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (nums[i] >= pivot)
        swap(nums, nextSwapped++, i);
    swap(nums, nextSwapped, r);

    final int count = nextSwapped - l + 1; // # of nums >= pivot
    if (count == k)
      return nums[nextSwapped];
    if (count > k)
      return quickSelect(nums, l, nextSwapped - 1, k);
    return quickSelect(nums, nextSwapped + 1, r, k - count);
  }

  private void swap(int[] nums, int i, int j) {
    final int temp = nums[i];
    nums[i] = nums[j];
    nums[j] = temp;
  }
}

Python

 
class Solution:
  def findKthLargest(self, nums: List[int], k: int) -> int:
    def quickSelect(l: int, r: int, k: int) -> int:
      pivot = nums[r]

      nextSwapped = l
      for i in range(l, r):
        if nums[i] >= pivot:
          nums[nextSwapped], nums[i] = nums[i], nums[nextSwapped]
          nextSwapped += 1
      nums[nextSwapped], nums[r] = nums[r], nums[nextSwapped]

      count = nextSwapped - l + 1  # number of nums >= pivot
      if count == k:
        return nums[nextSwapped]
      if count > k:
        return quickSelect(l, nextSwapped - 1, k)
      return quickSelect(nextSwapped + 1, r, k - count)

    return quickSelect(0, len(nums) - 1, k)

Approach 3: Quick Select with random pivot

  • Time:O(n) (average)
  • Space:O(1)

C++

 
class Solution {
 public:
  int findKthLargest(vector<int>& nums, int k) {
    return quickSelect(nums, 0, nums.size() - 1, k);
  }

 private:
  int quickSelect(vector<int>& nums, int l, int r, int k) {
    const int randIndex = rand() % (r - l + 1) + l;
    swap(nums[randIndex], nums[r]);
    const int pivot = nums[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (nums[i] >= pivot)
        swap(nums[nextSwapped++], nums[i]);
    swap(nums[nextSwapped], nums[r]);

    const int count = nextSwapped - l + 1;  // # of nums >= pivot
    if (count == k)
      return nums[nextSwapped];
    if (count > k)
      return quickSelect(nums, l, nextSwapped - 1, k);
    return quickSelect(nums, nextSwapped + 1, r, k - count);
  }
};

JAVA

 
class Solution {
  public int findKthLargest(int[] nums, int k) {
    return quickSelect(nums, 0, nums.length - 1, k);
  }

  private int quickSelect(int[] nums, int l, int r, int k) {
    final int randIndex = new Random().nextInt(r - l + 1) + l;
    swap(nums, randIndex, r);
    final int pivot = nums[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (nums[i] >= pivot)
        swap(nums, nextSwapped++, i);
    swap(nums, nextSwapped, r);

    final int count = nextSwapped - l + 1; // # of nums >= pivot
    if (count == k)
      return nums[nextSwapped];
    if (count > k)
      return quickSelect(nums, l, nextSwapped - 1, k);
    return quickSelect(nums, nextSwapped + 1, r, k - count);
  }

  private void swap(int[] nums, int i, int j) {
    final int temp = nums[i];
    nums[i] = nums[j];
    nums[j] = temp;
  }
}

Python

 
class Solution:
  def findKthLargest(self, nums: List[int], k: int) -> int:
    def quickSelect(l: int, r: int, k: int) -> int:
      randIndex = randint(0, r - l) + l
      nums[randIndex], nums[r] = nums[r], nums[randIndex]
      pivot = nums[r]

      nextSwapped = l
      for i in range(l, r):
        if nums[i] >= pivot:
          nums[nextSwapped], nums[i] = nums[i], nums[nextSwapped]
          nextSwapped += 1
      nums[nextSwapped], nums[r] = nums[r], nums[nextSwapped]

      count = nextSwapped - l + 1  # number of nums >= pivot
      if count == k:
        return nums[nextSwapped]
      if count > k:
        return quickSelect(l, nextSwapped - 1, k)
      return quickSelect(nextSwapped + 1, r, k - count)

    return quickSelect(0, len(nums) - 1, k)