Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 K개의 정렬된 리스트에서 최소 범위 찾는 방법

이번 글에서는 K개의 정렬된 정수 리스트가 주어졌을 때, 각 리스트에서 최소 하나 이상의 숫자를 포함하는 가장 작은 범위(smallest range)를 찾는 방법을 다룹니다.

여기서 두 범위를 비교할 때, 범위 [a, b]가 범위 [c, d]보다 작다는 것은 b - a < d - c 이거나, 크기가 같을 경우(b - a == d - c) a < c 인 경우를 의미합니다.

문제 예시

입력이 다음과 같다고 가정해 보겠습니다.

[[4,10,15,25,26], [0,9,14,20], [5,18,24,30]]

이 경우 출력은 [14, 18]이 됩니다. 세 리스트 모두 14와 18 사이의 범위에 속하는 숫자(첫 번째 리스트의 15, 두 번째 리스트의 14, 세 번째 리스트의 18)를 가지며, 이보다 더 좁은 범위는 존재하지 않기 때문입니다.

풀이 접근 방식

이 문제는 최소 힙(min-heap), 즉 우선순위 큐(priority queue)를 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • 각 리스트의 첫 번째 원소를 힙에 넣고, 현재까지의 최댓값(tempMaxRange)을 기록합니다.
  • 힙에서 최솟값을 꺼낼 때마다 (최댓값 - 최솟값)으로 범위 크기를 계산하여, 지금까지 찾은 최소 범위보다 작으면 갱신합니다.
  • 최솟값이 나온 리스트의 다음 원소를 힙에 삽입하고, 필요하면 최댓값을 갱신합니다.
  • 특정 리스트의 모든 원소를 소진하면 탐색을 종료합니다. 이 시점 이후에는 모든 리스트를 포함하는 범위를 만들 수 없기 때문입니다.

알고리즘 단계

  • minRange := 무한대, maxRange := 음의 무한대, rangeSize := 무한대, tempMinRange := 무한대, tempMaxRange := 음의 무한대로 초기화
  • n := 리스트의 개수(nums의 크기)
  • 크기가 n인 포인터 배열 pointers 생성 (각 리스트에서 현재 가리키는 위치 저장)
  • 우선순위 큐 pq 생성
  • i := 0부터 n 미만까지 반복:
    • { nums[i][0], i } 쌍을 pq에 삽입
    • tempMaxRange := max(tempMaxRange, nums[i][0])
  • 무한 루프 실행:
    • pq의 top 요소를 temp로 가져온 뒤 pop
    • tempMinRange := temp.first (현재 최솟값)
    • idx := temp.second (최솟값이 속한 리스트의 인덱스)
    • 만약 tempMaxRange - tempMinRange < rangeSize 라면:
      • rangeSize := tempMaxRange - tempMinRange
      • minRange := tempMinRange
      • maxRange := tempMaxRange
    • pointers[idx]를 1 증가
    • pointers[idx]가 nums[idx]의 크기와 같으면 루프 종료
    • 그렇지 않으면:
      • tempMaxRange := max(tempMaxRange, nums[idx][pointers[idx]])
      • { nums[idx][pointers[idx]], idx }를 pq에 삽입
  • 크기 2의 배열 ans 생성 후 ans[0] := minRange, ans[1] := maxRange 대입
  • ans 반환

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;
}
struct Comparator{
   bool operator() (pair <int, int> a, pair <int, int> b){
      return !(a.first < b.first);
   }
};
class Solution {
public:
   vector<int> smallestRange(vector<vector<int>>& nums) {
      int minRange = INT_MAX;
      int maxRange = INT_MIN;
      int rangeSize = INT_MAX;
      int tempMinRange, tempMaxRange, tempRangeSize;
      tempMinRange = INT_MAX;
      tempMaxRange = INT_MIN;
      int n = nums.size();
      vector <int> pointers(n);
      priority_queue < pair <int, int>, vector < pair <int, int> >, Comparator > pq;
      for(int i = 0; i < n; i++){
         pq.push({nums[i][0], i});
         tempMaxRange = max(tempMaxRange, nums[i][0]);
      }
      while(1){
         pair <int, int> temp = pq.top();
         pq.pop();
         tempMinRange = temp.first;
         int idx = temp.second;
         if(tempMaxRange - tempMinRange < rangeSize){
            rangeSize = tempMaxRange - tempMinRange;
            minRange = tempMinRange;
            maxRange = tempMaxRange;
         }
         pointers[idx]++;
         if(pointers[idx] == nums[idx].size())break;
         else{
            tempMaxRange = max(tempMaxRange, nums[idx][pointers[idx]]);
            pq.push({nums[idx][pointers[idx]], idx});
         }
      }
      vector <int> ans(2);
      ans[0] = minRange;
      ans[1] = maxRange;
      return ans;
   }
};
main(){
   Solution ob;
   vector<vector<int>> v = {{4,10,15,25,26},{0,9,14,20},{5,18,24,30}};
   print_vector(ob.smallestRange(v));
}

입력

{{4,10,15,25,26},{0,9,14,20},{5,18,24,30}};

출력

[14, 18]

시간 복잡도 분석

각 리스트의 원소는 힙에 최대 한 번씩 삽입되고 제거되므로, 전체 시간 복잡도는 O(N log K)입니다. 여기서 N은 모든 리스트의 원소 개수의 합이고, K는 리스트의 개수입니다. 힙에는 항상 최대 K개의 원소만 유지되기 때문에 로그 계수는 K에 비례합니다.

마무리

이 알고리즘은 우선순위 큐를 사용해 각 리스트의 포인터를 앞으로 이동시켜 가면서 가능한 모든 '후보 범위'를 효율적으로 탐색합니다. 슬라이딩 윈도우나 병합 정렬형 접근과 유사한 패턴으로, k-way 병합(k-way merge) 기법을 응용한 대표적인 문제이므로 코딩 인터뷰 준비에도 매우 유용합니다.