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

C++로 행별 정렬 행렬에서 중앙값(Median) 효율적으로 찾기

문제 소개

이 문제에서는 각 행이 오름차순으로 정렬되어 있는 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번째 작은 원소 찾기" 유형의 문제에도 널리 응용되므로 숙지해 두면 좋습니다.