[Leetcode] 658. Find K Closest Elements

문제

Given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.
An integer a is closer to x than an integer b if:
|a - x| < |b - x|, or|a - x| == |b - x| and a < b
 
Example 1:
Input: arr = [1,2,3,4,5], k = 4, x = 3Output: [1,2,3,4]
Example 2:
Input: arr = [1,1,2,3,4,5], k = 4, x = -1Output: [1,1,2,3]
 
Constraints:
1 <= k <= arr.length1 <= arr.length <= 104arr is sorted in ascending order.-104 <= arr[i], x <= 104

 

풀이1

처음 이 문제를 봤을 때, 각 element의 거리를 구해서, 거리를 기준으로 sort하고, k개의 elements 만 넣은 뒤, 다시 element를 기준으로 정렬하면 된다고 생각했다. 이 풀이도 물론 pass하지만 이렇게 풀었을 경우 quick sort 두번과, arrSized array & k-sized array 총 2개의 array를 만들어야 함으로 별로 efficient한 풀이는 아니다.

 

코드1 (passed, but less efficient)

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
#include <math.h>
typedef struct Node {
    int dist;
    int num;
} Node;

int cmp_dist(const void* a, const void* b) {
    return ((Node*)a)->dist - ((Node*)b)->dist;
}
int cmp(const void* a, const void* b) {
    return *((int*)a)- *((int*)b);
}
int* findClosestElements(int* arr, int arrSize, int k, int x, int* returnSize) {
    // find an array of dist from x
    Node* nodes = (Node*)malloc(arrSize * sizeof(Node));
    *returnSize = 0;
    for(int i=0; i<arrSize; i++) {
        nodes[i].dist = abs(arr[i] - x);
        nodes[i].num = arr[i];
        printf("%d, %d\n",nodes[i].dist, nodes[i].num );
    }
    // sort in dist and get k elements
    qsort(nodes, k, sizeof(Node), cmp_dist);
    // re-sort based on org element
    int* result = (int*)malloc(k * sizeof(int));
    for(int i=0; i<k; i++) {
        result[i] = nodes[i].num;
    }
    free(nodes);
    *returnSize = k;
    qsort(result, k, sizeof(int), cmp);

    return result;
}

 

풀이2

binarysearch가 어떤 한 condition에 가장 먼저 만족하는 index를 찾으므로, binary search를 활용하면 더 time efficient & space efficient 한 코드를 구할 수 있다. 

 

condition은 이미 문제에서 준 condition을 이용하면 된다.

|a - x| < |b - x|, or|a - x| == |b - x| and a < b

 

여기서 a는 arr[mid], b는 arr[mid+k]일 때, {arr[mid]와 x의 거리}가 {arr[mid+k]와 x의 거리}보다 같거나 작을 때의 첫 경우를 찾으려면 밑의 코드를 적으면 된다.

 

코드2

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
int* findClosestElements(int* arr, int arrSize, int k, int x, int* returnSize) {
    int left = 0, right = arrSize - k;

    while(left < right) {
        int mid = left + (right-left)/2;
        if(x - arr[mid] > (arr[mid+k] - x)) left = mid+1;
        else right = mid;
    }

    int* result = (int*)malloc(k*sizeof(int));
    for(int i=0;i<k;i++) {
        result[i] = arr[left+i];
    }
    (*returnSize) = k;
    return result;
}

 

코멘트:

#include <library.h>를 추가할때마다 leetcode time이 조금 더 걸린다. 최소한의 라이브러리만 추가해야겠다.