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

C++로 정렬된 두 배열의 중앙값 찾기 — 이진 탐색 기반 효율적 풀이

문제 개요

정렬된 두 개의 리스트가 주어졌을 때, 이 두 리스트를 합쳤을 때의 중앙값(median)을 구하는 것이 목표입니다. 예를 들어 배열이 [1,5,8][2,3,6,9]라면, 두 배열을 병합하면 [1,2,3,5,6,8,9]가 되고 전체 길이는 7(홀수)이므로 중앙값은 5입니다.

두 배열을 단순히 병합한 뒤 중앙값을 찾는 방법은 O(n+m)의 시간이 걸리지만, 이진 탐색(Binary Search)을 활용하면 O(log(min(n, m))) 시간 안에 훨씬 효율적으로 해결할 수 있습니다.

알고리즘 접근 방법

핵심 아이디어는 두 배열을 적절한 위치에서 '분할(partition)'하여, 왼쪽 부분의 모든 원소가 오른쪽 부분의 모든 원소보다 작거나 같도록 만드는 것입니다. 다음 단계로 진행합니다.

  • solve() 함수를 정의합니다. 이 함수는 배열 nums1nums2를 매개변수로 받습니다.
  • 만약 nums1의 크기가 nums2보다 크다면, solve(nums2, nums1)을 호출하여 항상 더 작은 배열을 기준으로 탐색하도록 합니다.
  • x := nums1의 크기, y := nums2의 크기
  • low := 0, high := x
  • totalLength := x + y
  • low <= high인 동안 다음을 반복합니다:
    • partitionX := low + (high - low) / 2
    • partitionY := (totalLength + 1) / 2 - partitionX
    • maxLeftX := partitionX가 0이면 -∞(INT_MIN), 그렇지 않으면 nums1[partitionX - 1]
    • minRightX := partitionX가 x이면 +∞(INT_MAX), 그렇지 않으면 nums1[partitionX]
    • maxLeftY := partitionY가 0이면 -∞(INT_MIN), 그렇지 않으면 nums2[partitionY - 1]
    • minRightY := partitionY가 y이면 +∞(INT_MAX), 그렇지 않으면 nums2[partitionY]
    • 만약 maxLeftX <= minRightY이고 maxLeftY <= minRightX라면 올바른 분할을 찾은 것입니다:
      • totalLength가 짝수이면 → ((maxLeftX와 maxLeftY 중 최댓값) + (minRightX와 minRightY 중 최솟값)) / 2를 반환
      • totalLength가 홀수이면 → maxLeftX와 maxLeftY 중 최댓값을 반환
    • 그렇지 않고 maxLeftX > minRightY라면 → high := partitionX - 1 (분할점을 왼쪽으로 이동)
    • 그 외의 경우 → low := partitionX + 1 (분할점을 오른쪽으로 이동)
  • 반복이 끝나면 0을 반환합니다.

C++ 구현 예제

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

class Solution {
public:
   double solve(vector<int>& nums1, vector<int>& nums2) {
      if(nums1.size() > nums2.size())
         return solve(nums2, nums1);
      int x = nums1.size();
      int y = nums2.size();
      int low = 0;
      int high = x;
      int totalLength = x + y;
      while(low <= high){
         int partitionX = low + (high - low) / 2;
         int partitionY = (totalLength + 1) / 2 - partitionX;

         int maxLeftX = (partitionX == 0 ? INT_MIN : nums1[partitionX - 1]);
         int minRightX = (partitionX == x ? INT_MAX : nums1[partitionX]);
         int maxLeftY = (partitionY == 0 ? INT_MIN : nums2[partitionY - 1]);
         int minRightY = (partitionY == y ? INT_MAX : nums2[partitionY]);

         if(maxLeftX <= minRightY && maxLeftY <= minRightX){
            if(totalLength % 2 == 0){
               return ((double)max(maxLeftX, maxLeftY) + (double)min(minRightX, minRightY)) / 2;
            }
            else{
               return max(maxLeftX, maxLeftY);
            }
         }
         else if(maxLeftX > minRightY)
            high = partitionX - 1;
         else
            low = partitionX + 1;
      }
      return 0;
   }
};

int main(){
   Solution ob;
   vector<int> v1 = {1,5,8}, v2 = {2,3,6,9};
   cout << ob.solve(v1, v2);
}

입력

[1,5,8], [2,3,6,9]

출력

5

동작 원리 요약

위 코드는 작은 배열을 기준으로 분할 위치를 이진 탐색으로 조정해 나갑니다. 각 단계에서 왼쪽 파티션의 최댓값(maxLeftX, maxLeftY)이 오른쪽 파티션의 최솟값(minRightX, minRightY)보다 작거나 같아지는 지점을 찾으면, 그 지점이 바로 두 배열이 병합되었을 때의 중앙 경계가 됩니다. 전체 길이가 짝수일 경우 중앙 경계 양쪽 값의 평균을, 홀수일 경우 왼쪽 파티션의 최댓값을 중앙값으로 반환합니다. 이 방식은 배열을 실제로 병합하지 않으므로 시간 복잡도 O(log(min(n, m))), 공간 복잡도 O(1)로 매우 효율적입니다.