Leetcode - 703. Kth Largest Element in a Stream

문제

주어진 배열에 값을 add할때마다 배열에서 k번째로 큰 값을 리턴하라.

Input
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output
[null, 4, 5, 5, 8, 8]

https://leetcode.com/problems/kth-largest-element-in-a-stream/

해결

k번째 큰값이라고 해서 꼭 max heap만 사용해야하는건 아님. 이 문제에선 min heap을 써야 더 수월하게 풀림.

배열을 min heap으로 만든 뒤, heap_size를 k가 될때 까지 pop하면 heap의 0번 index에는 k번째로 큰 값이 남아있게 됨. 이 성질 이용하여 해결하면 된다. 처음에는 max heap으로 하려고 했는데, 그러면 첫 pop할때만 정상적인 값을 얻을수 있고, 그 뒤 값이 제거가 되기 때문에 또다시 push를 해주거나 해야함. 그런데 min heap으로 만들고 위의 성질을 이용아래와 같이 수행하면 간단하게 해결됨.

우선 배열을 k개까지만 min heap으로 유지(단, 가장 큰값만 k개 남아있는 min heap임)
추가로 값이 add될 때마다 다음동작을 수행하면 heap[0]은 항상 k번째로 큰 값임.

• 값이 heap[0]보다 크면 heap[0]을 그 값으로 바꾸고 0부터 heapify_down

int kthLargestAdd(KthLargest* obj, int val) { 
    if (obj->heap_size < kth) {
        heap_push(obj, val);
    } else if (val > obj->heap[0]) {
        obj->heap[0] = val;
        heapify_down(obj->heap, 0, obj->heap_size);
    }
    return obj->heap[0];
}

값을kthLargestAdd() 을 통해 add하면 위의 성질을 유지하면서 k개의 heap이 유지된다. 재미있는 문제였다!

typedef struct {
    int *heap;
    int heap_size;
} KthLargest;
int kth;

void swap(int arr[], int a, int b)
{
    int temp = arr[a];
    arr[a] = arr[b];
    arr[b] = temp;
}
void heapify_up(int arr[], int curr_idx)
{
    int p_idx = (curr_idx - 1) / 2;
    
    if (curr_idx < 1)
        return;
    if (arr[curr_idx] < arr[p_idx]) {
        swap(arr, curr_idx, p_idx);
        heapify_up(arr, p_idx);
    }
}

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

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);
    }
}

int kthLargestAdd(KthLargest* obj, int val);

KthLargest* kthLargestCreate(int k, int* nums, int numsSize) {
    kth = k;
    KthLargest* q = (KthLargest *)malloc(sizeof(KthLargest));
    memset(q, 0, sizeof(KthLargest));
    q->heap = (int *)malloc(sizeof(int) * k);
    memset(q->heap, 0 , sizeof(int) * k);
    
    for (int i = 0; i < numsSize; i++)
        kthLargestAdd(q, nums[i]);
    return q;
}

int kthLargestAdd(KthLargest* obj, int val) { 
    if (obj->heap_size < kth) {
        heap_push(obj, val);
    } else if (val > obj->heap[0]) {
        obj->heap[0] = val;
        heapify_down(obj->heap, 0, obj->heap_size);
    }
    return obj->heap[0];
}

void kthLargestFree(KthLargest* obj) {
    free(obj->heap);    
    free(obj);    
}

/**
 * Your KthLargest struct will be instantiated and called as such:
 * KthLargest* obj = kthLargestCreate(k, nums, numsSize);
 * int param_1 = kthLargestAdd(obj, val);
 
 * kthLargestFree(obj);
*/