Leetcode - 973. K Closest Points to Origin

문제

좌표평면에 좌표들이 주어질때, 원점에서 가장 가까운 순서대로 k개 만큼 출력하기.

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]

https://leetcode.com/problems/k-closest-points-to-origin/

해결 O(N logN)

크기순으로 k개만큼만 출력하는 전형적인 heap(우선순위큐)문제. qsort등으로 전체를 정렬해서 풀면 O(N logN)이지만, heap을 이용해 구하면 O(K logN)에 해결이 가능하다. 입력된 좌표를 하나씩 heap_push 하고 O(N logN), K개 만큼 heap_pop하면 O(K logN) 해결. (개선포인트) 입력된 좌표를 heap_push하지 말고, 배열에 일단 O(N) 에 저장한 뒤, heap배열로 만들면 O(N)에 heap구조로 만들 수 있다.

#include <math.h>

struct point {
    double dist;
    int idx;
};
struct pq {
    int heap_size;
    struct point *heap;
};
struct pq *q;

void swap(struct point arr[], int a, int b)
{
    struct point tmp = arr[a];
    arr[a] = arr[b];
    arr[b] = tmp;
}
void heapify_up(struct point arr[], int curr_idx)
{
    int p_idx = (curr_idx - 1) / 2;
    if (curr_idx < 1)
        return;
    if (arr[curr_idx].dist < arr[p_idx].dist) {
        swap(arr, curr_idx, p_idx);
        heapify_up(arr, p_idx);
    }
}
void heap_push(struct point arr[], struct point p)
{
    arr[q->heap_size].dist = p.dist;
    arr[q->heap_size].idx = p.idx;
    q->heap_size++;
    heapify_up(q->heap, q->heap_size - 1);
}

void heapify_down(struct point 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].dist < arr[min_idx].dist)
        min_idx = l_idx;
    if (r_idx < size && arr[r_idx].dist < arr[min_idx].dist)
        min_idx = r_idx;
    if (min_idx != p_idx) {
        swap(arr, min_idx, p_idx);
        heapify_down(arr, min_idx, size);
    }
}
struct point heap_pop(struct point arr[])
{
    struct point ret = arr[0];
    q->heap_size--;
    arr[0] = q->heap[q->heap_size];
    heapify_down(arr, 0, q->heap_size);
    return ret;
}
void print_heap(struct point arr[])
{
    for (int i = 0; i < q->heap_size; i++)
        printf("(%f,%d)", arr[i].dist, arr[i].idx);
    printf("\n");
}
/**
 * Return an array of arrays of size *returnSize.
 * The sizes of the arrays are returned as *returnColumnSizes array.
 * Note: Both returned array and *columnSizes array must be malloced, assume caller calls free().
 */
int** kClosest(int** points, int pointsSize, int* pointsColSize, int k, int* returnSize, int** returnColumnSizes){
    q = (struct pq *)malloc(sizeof(struct pq));
    memset(q, 0, sizeof(struct pq));
    q->heap = (struct point *)malloc(sizeof(struct point) * pointsSize);
    memset(q->heap, 0, sizeof(struct point) * pointsSize);
    int **ret = (int **)malloc(sizeof(int *) * k);
    *returnSize = k;
    *returnColumnSizes = (int*)malloc(sizeof(int) * k);
    struct point tmp;
    
    for (int i = 0; i < pointsSize; i++) {
        tmp.dist = sqrt(points[i][0] * points[i][0] + points[i][1] * points[i][1]);
        tmp.idx = i;
        heap_push(q->heap, tmp);
    }
    //print_heap(q->heap);
    for (int i = 0; i < k; i++) {
        ret[i] = (int *)malloc(sizeof(int) * 2);
        (*returnColumnSizes)[i] = 2;
        tmp = heap_pop(q->heap);
        ret[i][0] = points[tmp.idx][0];
        ret[i][1] = points[tmp.idx][1];
    }
    return ret;
}