C++でソート済みの2つの配列の中央値を求める方法【二分探索】

問題概要

2つのソート済み配列が与えられ、それらを合わせた全体の中央値を求めます。例えば、配列が [1,5,8] と [2,3,6,9] の場合、マージすると [1,2,3,5,6,8,9] となり、中央の要素である 5 が答えになります。

この問題は、単純に2つの配列をマージしてから中央値を取り出す方法（計算量 O(m+n)）でも解けますが、二分探索を活用することで、O(log(min(m, n))) まで計算量を抑えることができます。ポイントは「短い方の配列」だけを対象に二分探索を行い、両配列を左右半分に分ける適切な分割位置（パーティション）を見つけることです。

アルゴリズムの手順

1. 関数 findMedianSortedArrays を定義し、配列 nums1 と nums2 を受け取ります。
2. nums1 のサイズが nums2 より大きい場合は、findMedianSortedArrays(nums2, nums1) を呼び出して引数を入れ替えます（常に短い方の配列を基準に探索するため）。
3. x := nums1 のサイズ、y := nums2 のサイズとします。
4. low := 0、high := x、totalLength := x + y と初期化します。
5. low <= high の間、次の処理を繰り返します。
   • partitionX := low + (high − low) / 2
   • partitionY := (totalLength + 1) / 2 − partitionX
   • maxLeftX = partitionX が 0 の場合は −∞、それ以外は nums1[partitionX − 1]
   • minRightX = partitionX が x の場合は +∞、それ以外は nums1[partitionX]
   • maxLeftY = partitionY が 0 の場合は −∞、それ以外は nums2[partitionY − 1]
   • minRightY = partitionY が y の場合は +∞、それ以外は nums2[partitionY]
6. 条件に応じて判定を行います。
   • maxLeftX ≤ minRightY かつ maxLeftY ≤ minRightX の場合：正しい分割が見つかった状態です。totalLength が偶数なら (max(maxLeftX, maxLeftY) + min(minRightX, minRightY)) / 2 を返し、奇数なら max(maxLeftX, maxLeftY) を返します。
   • maxLeftX > minRightY の場合：分割位置が大きすぎるため、high := partitionX − 1 として左側へ範囲を狭めます。
   • それ以外の場合：low := partitionX + 1 として右側へ範囲を狭めます。
7. ループを抜けた場合は 0 を返します。

C++での実装例

理解を深めるために、実際のC++コードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        if(nums1.size()>nums2.size())
            return findMedianSortedArrays(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.findMedianSortedArrays(v1, v2);
}

入力

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

出力

5

計算量の評価

このアルゴリズムでは、常に短い方の配列に対して二分探索を行うため、時間計算量は O(log(min(m, n))) となります。また、追加の配列を作成せずインデックスの比較だけで処理が完結するため、空間計算量は O(1) です。要素数が多い大規模なデータセットでも高速に動作する、非常に効率的な解法と言えます。