C++로 k개의 정렬된 배열에서 m번째로 작은 값 찾는 방법

이 문제에서는 크기가 서로 다른 k개의 정렬된 배열이 주어지며, 이 배열들을 합쳤을 때의 m번째로 작은 값을 찾아야 합니다.

문제 설명

k개의 정렬된 배열을 하나로 병합한 배열을 만들었을 때, 그 배열에서 m번째로 작은 원소를 구하는 것이 목표입니다.

예제로 문제 이해하기

입력:

m = 4
arr[][] = { {4, 7},
            {2, 5, 6},
            {3, 9, 12, 15, 19} }

출력: 5

설명: 모든 배열을 병합하고 정렬하면 다음과 같습니다.

2, 3, 4, 5, 6, 7, 9, 12, 15, 19

이 정렬된 배열에서 4번째 원소는 5입니다.

해결 접근 방법

1. 단순한 방법 — 병합 후 정렬

가장 간단한 방법은 모든 배열의 원소를 하나의 배열에 병합한 뒤 오름차순으로 정렬하는 것입니다. 정렬이 완료되면 인덱스 (m-1) 위치의 값이 곧 m번째로 작은 원소가 됩니다.

다만 이 방법은 전체 원소를 모두 저장하고 정렬해야 하므로 시간 복잡도와 공간 복잡도 측면에서 비효율적입니다.

2. 효율적인 방법 — 최소 힙(Min Heap) 활용

더 효율적인 해결책은 최소 힙(Min Heap) 자료구조를 사용하는 것입니다. 동작 과정은 다음과 같습니다.

각 배열의 첫 번째 원소를 최소 힙에 삽입합니다. 이때 어떤 배열의 몇 번째 원소인지를 함께 저장합니다.
힙에서 가장 작은 원소를 꺼내고(pop), 해당 원소가 속한 배열의 다음 원소를 힙에 삽입합니다.
위 과정을 m번 반복합니다.
m번째로 꺼낸 값이 바로 우리가 찾는 답입니다.

이 방식은 실제로 배열을 병합하지 않고도 필요한 만큼(m개)만 처리하므로 훨씬 효율적입니다. 시간 복잡도는 O(m log k)가 됩니다.

C++ 구현 예제

#include <bits/stdc++.h>
using namespace std;

typedef pair<int, pair<int, int> > ppi;

int findMSmallestElement(vector<vector<int> > sortedArr, int m) {
    priority_queue<ppi, vector<ppi>, greater<ppi> > priorQueue;

    // 각 배열의 첫 번째 원소를 힙에 삽입
    for (int i = 0; i < sortedArr.size(); i++)
        priorQueue.push({ sortedArr[i][0], { i, 0 } });

    int count = 0;
    int i = 0, j = 0;
    while (count < m && priorQueue.empty() == false) {
        ppi curr = priorQueue.top();
        priorQueue.pop();
        i = curr.second.first;
        j = curr.second.second;
        // 같은 배열의 다음 원소를 힙에 추가
        if (j + 1 < sortedArr[i].size())
            priorQueue.push({ sortedArr[i][j + 1], { i, (j + 1) } });
        count++;
    }
    return sortedArr[i][j];
}

int main() {
    vector<vector<int> > arr{ {4 , 7},
                              {2, 5, 6},
                              {3, 9, 12, 15, 19}};
    int m = 6;
    cout<<m<<"th smallest value in k sorted arrays is "<<findMSmallestElement(arr, m);

    return 0;
}

실행 결과

6th smallest value in k sorted arrays is 7

마무리

k개의 정렬된 배열에서 m번째로 작은 값을 구할 때, 단순 병합·정렬 방식보다는 최소 힙을 활용한 방법이 효율성 면에서 큰 이점을 제공합니다. 특히 배열의 개수나 원소 개수가 많을 때 이 방법의 장점이 더욱 두드러집니다.