Leetcode - 378. Kth Smallest Element in a Sorted Matrix

문제

n x n크기의 matrix가 주어진다. 각각의 row와 column은 오름차순 정렬되어있다고 할때 k번째로 작은 값을 구하라. 단 공간복잡도는 O(n^2) 미만으로 해결하라.

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

해결 O(n^2 logn) / O(k)

k번째라는 말이 붙으면 일단 heap으로 푸는 문제를 떠올려야함. heap_add()함수를 주의깊게 볼것. max heap으로 k개까지만 push하면 heap의 0번 인덱스는 k번째로 작은 값이 된다.

유사한 풀이로는

• 215. Kth Largest Element in an Array
• 703. Kth Largest Element in a Stream

struct pq{
    int size;
    int *arr;
};
struct pq *q;

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 largest = p_idx;
    int l_idx = p_idx * 2 + 1;
    int r_idx = p_idx * 2 + 2;
    
    if (l_idx < size && arr[l_idx] > arr[largest])
        largest = l_idx;
    if (r_idx < size && arr[r_idx] > arr[largest])
        largest = r_idx;
    if (largest != p_idx) {
        swap(arr, largest, p_idx);
        heapify_down(arr, largest, size);
    }
}

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

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

void heap_add(struct pq *q, int val, int kth)
{
    if (q->size < kth) {
        heap_push(q, val);
    } else if (val < q->arr[0]) {
        q->arr[0] = val;
        heapify_down(q->arr, 0, q->size);
    }
}

int kthSmallest(int** matrix, int matrixSize, int* matrixColSize, int k){
    int ret = 0;
    int hsize = k;
    q = (struct pq *)malloc(sizeof(struct pq));
    q->arr = (int *)malloc(hsize * sizeof(int));
    q->size = 0;
    
    for (int i = 0; i < matrixSize; i++)
        for (int j = 0; j < *matrixColSize; j++)
            heap_add(q, matrix[i][j], k); // Max Heap
    return q->arr[0];
}