정렬된 두 개의 배열 A1과 A2, 그리고 값 k가 주어졌다고 가정해 봅시다. 우리는 A1의 한 원소와 A2의 한 원소로 구성되는 쌍(pair) (u, v)을 정의해야 합니다. 그리고 합이 가장 작은 k개의 쌍, 즉 [(u1, v1), (u2, v2), …, (uk, vk)]을 찾아야 합니다.
예를 들어 A1 = [1, 7, 11], A2 = [2, 4, 6], k = 3이라면 출력 결과는 [(1, 2), (1, 4), (1, 6)]이 됩니다.
문제 해결 접근 방식
이 문제는 최소 힙(min-heap) 역할을 하는 우선순위 큐(priority queue)를 활용하면 효율적으로 해결할 수 있습니다. 알고리즘은 다음과 같이 진행됩니다.
- 두 값 a, b와 인덱스(index)를 저장하는 데이터 구조체를 하나 정의합니다.
- 합이 작은 순서대로 정렬되는 우선순위 큐를 생성합니다.
- n := A1의 크기, m := A2의 크기로 설정합니다.
- n 또는 m이 0이면 빈 결과를 반환합니다.
- 결과를 저장할 행렬 ret을 생성합니다.
- i가 0부터 n-1까지일 때, 각 i에 대해 (A1[i], A2[0], 0) 데이터를 큐에 삽입합니다.
- 큐가 비어 있지 않고 k가 0이 아닌 동안 다음을 반복합니다.
- curr := 큐의 최상단 요소를 꺼내고 해당 요소를 삭제합니다.
- curr의 first_val과 second_val을 ret에 삽입합니다.
- 만약 curr의 인덱스 + 1 < m이라면, (first_val, A2[인덱스 + 1], 인덱스 + 1) 데이터를 큐에 삽입합니다.
- k를 1씩 감소시킵니다.
- 최종적으로 ret을 반환합니다.
핵심 아이디어는 각 원소 u ∈ A1에 대해 A2의 첫 번째 원소와 짝을 지어 큐에 넣고, 쌍을 하나 꺼낼 때마다 같은 u와 A2의 다음 원소로 만들 수 있는 후보 쌍을 다시 큐에 넣는 것입니다. 이렇게 하면 항상 현재 가능한 가장 작은 합을 가진 쌍부터 순서대로 얻을 수 있습니다.
예제 코드
다음 구현을 통해 더 잘 이해할 수 있습니다.
#include <bits/stdc++.h>
#include <stack>
using namespace std;
struct Data{
int firstVal, secondVal, idx;
Data(int a, int b, int c){
firstVal = a;
secondVal = b;
idx = c;
}
};
struct Comparator{
bool operator()(Data a, Data b){
return !(a.firstVal + a.secondVal < b.firstVal + b.secondVal);
}
};
class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
priority_queue <Data, vector <Data>, Comparator> pq;
int n = nums1.size();
int m = nums2.size();
if(!n || !m)return {};
vector < vector <int> > ret;
for(int i = 0; i < n; i++){
pq.push(Data(nums1[i], nums2[0], 0));
}
while(!pq.empty() && k--){
Data curr = pq.top();
pq.pop();
ret.push_back({curr.firstVal, curr.secondVal});
if(curr.idx + 1 < m){
pq.push(Data(curr.firstVal, nums2[curr.idx + 1], curr.idx + 1));
}
}
return ret;
}
};
void print(vector <int> const &arr) {
cout<<"[";
for(int i=0; i < arr.size(); i++)
std::cout << arr.at(i) <<",";
cout<<"]";
}
int main() {
vector<int> nums1{1,7,11};
vector<int> nums2{2,4,6};
int k = 3;
Solution ob1;
vector<vector<int>> numsRet;
numsRet = ob1.kSmallestPairs(nums1, nums2, k);
cout<<"[";
for (vector<int> x : numsRet) {
print(x);
cout<<",";
}
cout<<"]"<<endl;
return 0;
}입력
nums1 = [1,7,11] nums2 = [2,4,6] k = 3
출력
[[1,2],[1,4],[1,6]]
복잡도 분석
이 알고리즘의 시간 복잡도는 O(k log n)입니다. 초기화 단계에서 n개의 요소를 힙에 삽입하는 데 O(n log n)이 소요되며, 이후 k번 반복하면서 매번 힙 연산(log n)이 발생하기 때문입니다. 공간 복잡도는 우선순위 큐에 저장되는 요소 수에 비례하여 O(n + k)입니다. 완전 탐색으로 모든 n × m개의 쌍을 생성하는 O(n × m log(n × m)) 방식보다 훨씬 효율적입니다.