Leetcode - 215. Kth Largest Element in an Array

문제

주어진 배열에서 k번째로 큰 값을 구하기

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4

https://leetcode.com/problems/kth-largest-element-in-an-array/

해결 O(nlogn)

k번째로 크거나 작은값은 전형적인 heap문제이다. 정렬을 하면 O(nlogn)에 풀수 있는데, heap을 사용하면 O(klogn)만에 풀 수 있다.
여기서는 O(nlogn) 방법의 heap으로 풀었다. k번째 largest값을 구할때 k사이즈만큼 min heap을 만들면 heap[0]은 항상 k번째로 큰 값이 된다. heap_add()를 통해 이를 구현할 수 있다. 이제 이 풀이방법은 확실히 익힌것 같음.
지난번에 풀었던 유사 문제로는 -> Leetcode - 703. Kth Largest Element in a Stream

void heap_add(int val)
{
    if (q->heap_size < kth) {
        heap_push(q->heap, val);
    } else {
        if (val > q->heap[0]) {
            q->heap[0] = val;
            heapify_down(q->heap, 0, q->heap_size);
        }
    }
}

물론 이방법 말고, build_heap을 한 뒤O(N) -> heap_pop을 O(logN), k 번하면 최종 시간복잡도는 O(K logN)으로 풀 수 있다.

/* k size min heap[] -> heap[0] : kth max element */

struct pq {
    int heap_size;
    int *heap;
};
struct pq *q;
int kth;

void swap(int arr[], int a, int b)
{
    int tmp = arr[a];
    arr[a] = arr[b];
    arr[b] = tmp;
}

void heapify_down(int arr[], int p_idx, int size)
{
    int min_idx = p_idx;
    int l_idx = p_idx * 2 + 1;
    int r_idx = p_idx * 2 + 2;
    
    if (l_idx < size && arr[l_idx] < arr[min_idx])
        min_idx = l_idx;
    if (r_idx < size && arr[r_idx] < arr[min_idx])
        min_idx = r_idx;
    if (min_idx != p_idx) {
        swap(arr, min_idx, p_idx);
        heapify_down(arr, min_idx, size);
    }
}

void heapify_up(int arr[], int curr_idx)
{
    int p_idx = (curr_idx - 1) >> 1;
    
    if (curr_idx < 1)
        return;
    if (arr[curr_idx] < arr[p_idx]) {
        swap(arr, p_idx, curr_idx);
        heapify_up(arr, p_idx);
    }
}

void heap_push(int arr[], int val)
{
    arr[q->heap_size] = val;
    q->heap_size++;
    heapify_up(arr, q->heap_size - 1);
}

void heap_add(int val)
{
    if (q->heap_size < kth) {
        heap_push(q->heap, val);
    } else {
        if (val > q->heap[0]) {
            q->heap[0] = val;
            heapify_down(q->heap, 0, q->heap_size);
        }
    }
}

int findKthLargest(int* nums, int numsSize, int k){
    q = (struct pq *)calloc(1, sizeof(struct pq));
    q->heap = (int *)calloc(k, sizeof(int)); // heap size is maximum k
    kth = k;
    
    for (int i = 0; i < numsSize; i++)
        heap_add(nums[i]);
    return q->heap[0];
}