k-й наибольший элемент в массиве

Источник: https://leetcode.com/problems/kth-largest-element-in-an-array/

Дан массив целых чиселnums и целое число k, вернуть k -й наибольший элемент в массиве, как сли бы массив был отсортирован.

Задачу нужно решить за O(n).

Пример 1:

Дано: nums = [3,2,1,5,6,4], k = 2
Результат: 5

Пример 2:

Дано: nums = [3,2,3,1,2,4,5,5,6], k = 4
Результат: 4
 

Ограничения:

• 1 <= k <= nums.length <= 105
• -104 <= nums[i] <= 104

Идея решения 1

Будем использовать кучу с наименьшим элементом в вершине, ограничив размер кучи числом k. Добавляя все элементы в массиве, один за другим, как только мы превысим размер кучи, то удалим наименьший элемент. Таким образом, в куче будет всегда находиться  k наибольших элементов.

После обработки всего массива, k-й наибольший элемент будет находится в вершине кучи.

Сложность алгоритма:

1. добавление в кучу размера k занимает O(log k).
2. мы добавляем N элментов в кучу, где N - размер массива.
3. Итого общая сложность = O(N log k)

Имплементация 1

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // init heap 'the smallest element first'
        PriorityQueue<Integer> heap =
            new PriorityQueue<Integer>((n1, n2) -> n1 - n2);

        // keep k largest elements in the heap
        for (int n: nums) {
          heap.add(n);
          if (heap.size() > k)
            heap.poll();
        }

        // output
        return heap.poll();        
  }
}

Идея решения 2

Алгоритм QuickSelect используется для поиска k-го наименьшего элемента и его сложность = O(N) в среднем.

Имплементация 2

import java.util.Random;
class Solution {
  int [] nums;

// поменять a и b местами в алгоритме 
public void swap(int a, int b) {
    int tmp = this.nums[a];
    this.nums[a] = this.nums[b];
    this.nums[b] = tmp;
  }

public static int partition(int left, int right){

        int pivot_index = left;


        int pivot = this.nums[pivot_index]

        for (int i = left; i <= right; i++) {

            //всё, что меньше pivot'а идет влево, pivot_index корректируется


            if (this.nums[i] < pivot) {

                swap(pivot_index, i);

                pivot_index++;

            }

        }

  

        // перемещаем pivot на финальную позицию


        swap(right, pivot_index);


  

        return pivot_index;

    }

// QUICKSELECT
  public int quickselect(int left, int right, int k_smallest) {
   
    if (left == right) // мы рассматриваем интервал из одного элменета
      return this.nums[left];  // этот единственный элемент и возвращаем

    pivot_index = partition(left, right);

    // если pivot получился на (N - k)й позиции
    if (k_smallest == pivot_index)
      return this.nums[k_smallest]; // возвращаем элемент в pivot_index
    // если (N-k) левее, то уменьшаем pivot_index
    else if (k_smallest < pivot_index)
      return quickselect(left, pivot_index - 1, k_smallest);
    // иначе увеличиваем
    return quickselect(pivot_index + 1, right, k_smallest);
  }

// СТАРТОВАЯ ФУНКЦИЯ
  public int findKthLargest(int[] nums, int k) {
    this.nums = nums;
    int size = nums.length;
    // k-й наибольший - это (N - k)й наименьший
    return quickselect(0, size - 1, size - k);
  }
}