문제 소개
이 문제에서는 각 행이 오름차순으로 정렬되어 있는 2차원 배열 mat[r][c]가 주어지며, 우리가 할 일은 이 행렬의 중앙값(median)을 찾는 것입니다.
문제 설명 — 행렬의 모든 원소를 하나의 배열로 펼쳤을 때의 중앙값을 구해야 합니다.
예제로 문제 이해하기
입력
mat = {
{2, 4, 7},
{5, 6, 8},
{4, 8, 9}
}출력
6
설명
행렬의 원소들을 배열에 저장하면 다음과 같습니다.
{2, 4, 4, 5, 6, 7, 8, 8, 9}
중앙값은 6입니다.해결 접근 방법
1. 단순한 방법: 정렬 후 중앙값 추출
가장 직관적인 해결책은 행렬의 모든 원소를 별도의 배열에 저장한 뒤, 배열을 정렬하고 가운데 원소를 꺼내는 것입니다. 하지만 이 방법은 전체 원소 수가 많을 경우 O(r·c·log(r·c))의 정렬 비용이 들어 비효율적입니다.
2. 효율적인 방법: 이진 탐색(Binary Search) 활용
훨씬 효과적인 접근 방식은 "전체 원소 중 중앙값보다 작거나 같은 원소는 정확히 (r×c)/2개 존재한다"는 성질을 이용하는 것입니다. 즉, 이 조건을 만족하는 값을 찾으면 그것이 곧 중앙값입니다.
구체적인 과정은 다음과 같습니다.
- 행이 정렬되어 있으므로 최솟값은 첫 번째 열에서, 최댓값은 마지막 열에서 찾습니다.
- 최솟값과 최댓값을 양 끝으로 하는 범위에서 이진 탐색을 수행하며 범위의 중간값(mid)을 구합니다.
- mid보다 작거나 같은 원소의 개수를 센 뒤, 개수가 (r×c+1)/2보다 작으면 하한을 mid+1로 올리고, 그렇지 않으면 상한을 mid로 낮춥니다.
- 탐색 범위가 좁혀져 수렴하면 그 값이 바로 중앙값입니다.
mid 이하의 원소 개수를 셀 때는 각 행마다 mid보다 큰 첫 번째 원소의 위치를 찾아 누적하면 되는데, C++의 내장 함수인 upper_bound()를 사용하면 각 행을 O(log c) 만에 처리할 수 있습니다.
전체 시간 복잡도는 O(r·log(c)·log(max−min))로, 모든 원소를 정렬하는 방법보다 훨씬 빠릅니다.
구현 예제
#include<bits/stdc++.h>
using namespace std;
#define c 3
#define r 3
int findMedian(int mat[][c]) {
int smallest = INT_MAX, largest = INT_MIN;
for (int i=0; i<r; i++) {
if (mat[i][0] < smallest)
smallest = mat[i][0];
if (mat[i][c-1] > largest)
largest = mat[i][c-1];
}
while (smallest < largest){
int mid = smallest + (largest - smallest) / 2;
int smallCount = 0;
for (int i = 0; i < r; ++i)
smallCount += upper_bound(mat[i], mat[i]+c, mid) - mat[i];
if (smallCount < ((r * c + 1) / 2))
smallest = mid + 1;
else
largest = mid;
}
return smallest;
}
int main(){
int mat[][c]= { {2, 5, 7}, {4, 6, 8}, {1, 8, 9} };
cout<<"The median of the matrix is "<<findMedian(mat);
return 0;
}실행 결과
The median of the matrix is 6
마무리
행별로 정렬된 행렬의 중앙값을 구할 때는 모든 원소를 정렬하는 대신, 이진 탐색과 upper_bound()를 결합하면 훨씬 적은 연산으로 답을 구할 수 있습니다. 이 기법은 "k번째 작은 원소 찾기" 유형의 문제에도 널리 응용되므로 숙지해 두면 좋습니다.