C++로 정렬된 배열에서 x에 가장 가까운 k개의 요소 찾기

문제 개요

정렬된 배열과 두 개의 정수 k, x가 주어졌을 때, 배열에서 x에 가장 가까운 k개의 요소를 찾는 문제입니다. 결과는 반드시 오름차순으로 정렬되어야 하며, 거리가 같은 경우(동률)에는 항상 더 작은 값을 우선적으로 선택합니다.

예를 들어 입력 배열이 [1,2,3,4,5]이고 k = 4, x = 3이라면, 출력은 [1,2,3,4]가 됩니다.

접근 방법: 이진 탐색 활용

배열이 이미 정렬되어 있으므로, 이진 탐색(Binary Search)을 사용하면 O(log(n-k)) 시간 복잡도로 효율적으로 해결할 수 있습니다. 알고리즘은 다음과 같습니다.

• 결과를 담을 배열 ans를 선언합니다.
• low = 0, high = 배열 크기 - k로 초기화합니다.
• low < high인 동안 다음을 반복합니다.
  • mid = low + (high - low) / 2 로 중간 지점을 계산합니다.
  • x - arr[mid] > arr[mid + k] - x 이면 low = mid + 1로 갱신하고, 그렇지 않으면 high = mid로 갱신합니다.
• 인덱스 low부터 low + k까지의 요소를 ans에 삽입합니다.
• ans를 반환합니다.

핵심 아이디어는 'x에서 arr[mid]까지의 거리'와 'arr[mid + k]에서 x까지의 거리'를 비교하여, 윈도우를 왼쪽으로 이동할지 오른쪽으로 이동할지 결정하는 것입니다. 이 비교 조건은 자연스럽게 동률일 때 더 작은 요소를 우선하도록 처리해 줍니다.

C++ 구현 예제

아래 코드를 통해 실제 구현을 확인해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
cout << "]"<<endl;
}
class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        vector <int> ans;
        int low = 0;
        int high = arr.size() - k;
        while(low < high){
            int mid = low + (high - low)/2;
            if(x - arr[mid] > arr[mid + k] - x){
                low = mid + 1;
            }
            else high = mid;
        }
        for(int i = low ; i < low + k ; i++)ans.push_back(arr[i]);
        return ans;
    }
};
main(){
    Solution ob;
    vector<int> v = {1,2,3,4,5};
    print_vector(ob.findClosestElements(v, 4, 3));
}

입력

[1,2,3,4,5]
4
3

출력

[1,2,3,4]

정리

이 문제는 정렬된 배열의 특성을 활용한 이진 탐색 기법의 좋은 예입니다. 단순히 모든 요소와 x의 거리를 계산하는 방식(O(n log n))보다 훨씬 효율적이며, 특히 대용량 데이터에서 성능 차이가 크게 나타납니다. 시간 복잡도는 O(log(n-k) + k), 공간 복잡도는 O(k)입니다.