二つのソート済み配列の中央値の効率的な探索

問題概要

二つの昇順にソートされた整数配列 nums1nums2 が与えられます。それぞれの配列のサイズは mn です。これら二つの配列を結合した場合の中央値を求めてください。

このアルゴリズムの時間計算量は O(log (m+n)) である必要があります。


例1:
入力:nums1 = [1,3], nums2 = [2]
出力:2.00000
解説:結合配列 = [1,2,3] 、中央値 2

アプローチ1:マージによる方法 (時間計算量 O(m+n))

この問題に対する最も直接的なアプローチは、二つの与えられたソート済み配列 nums1nums2 を一つの大きなソート済み配列にマージすることです。マージが完了したら、その結合された配列から中央値を計算します。

  • マージ処理は、両配列の先頭から要素を順に比較し、小さい方を新しい配列に追加していくことで行われます。
  • 結合された配列の長さが奇数の場合、中央値は単純に配列の中央要素です。
  • 結合された配列の長さが偶数の場合、中央値は中央の二つの要素の平均となります。
  • この方法の時間計算量は O(m+n)、空間計算量も O(m+n) です。しかし、問題の要件である O(log (m+n)) を満たさないため、より効率的なアプローチが必要です。

class Solution {
    public double findMedianSortedArrays(int[] arrA, int[] arrB) {
        int lenA = arrA.length;
        int lenB = arrB.length;
        int totalLen = lenA + lenB;

        // マージ結果を格納する配列
        int[] mergedResult = new int[totalLen];
        int mergedIdx = 0;
        int ptrA = 0; // arrA の現在位置
        int ptrB = 0; // arrB の現在位置
        
        // arrA と arrB をソート順にマージ
        while (ptrA < lenA && ptrB < lenB) {
            if (arrA[ptrA] < arrB[ptrB]) {
                mergedResult[mergedIdx++] = arrA[ptrA++];
            } else {
                mergedResult[mergedIdx++] = arrB[ptrB++];
            }
        }
        
        // arrA の残りの要素を追加
        while (ptrA < lenA) {
            mergedResult[mergedIdx++] = arrA[ptrA++];
        }
        
        // arrB の残りの要素を追加
        while (ptrB < lenB) {
            mergedResult[mergedIdx++] = arrB[ptrB++];
        }
        
        // マージされた配列から中央値を計算
        if (totalLen % 2 == 1) {
            // 長さが奇数の場合
            return (double) mergedResult[totalLen / 2];
        } else {
            // 長さが偶数の場合
            return (double) (mergedResult[totalLen / 2 - 1] + mergedResult[totalLen / 2]) / 2.0;
        }
    }
}

アプローチ2:二分探索 (K番目に小さい要素の探索) (時間計算量 O(log(m+n)))

時間計算量の制約 O(log (m+n)) を満たすためには、二分探索の原理を応用する必要があります。この問題は、本質的に「二つのソート済み配列からK番目に小さい要素を見つける」という問題に帰着できます。

  • 結合された配列の長さが奇数の場合、中央値は (m+n)/2 + 1 番目に小さい要素です。
  • 結合された配列の長さが偶数の場合、中央値は (m+n)/2 番目と (m+n)/2 + 1 番目に小さい要素の平均です。

K番目に小さい要素を見つける再帰的ロジック

二つの配列 arr1arr2 と、それぞれの探索範囲(start1, end1, start2, end2)、そして目標とする順序 targetK を引数に取る再帰関数 getKthSmallest を設計します。

  1. 基本ケース:

    • もし arr1 の探索範囲が空であれば(currentLen1 == 0)、arr2 の中から targetK 番目に小さい要素(start2 + targetK - 1 インデックス)を返します。arr2 が空の場合も同様です。
    • もし targetK1 であれば、arr1[start1]arr2[start2] のうち小さい方が1番目に小さい要素なので、それを返します。
    • 再帰呼び出しを安定させるため、arr1 が常に短い方の配列になるように、必要に応じて配列を入れ替えて呼び出します。
  2. 探索範囲の縮小:

    • それぞれの配列から targetK / 2 個の要素を選択し、それらを中央値の候補として比較します。具体的には、arr1start1 + targetK / 2 - 1 番目の要素と、arr2start2 + targetK / 2 - 1 番目の要素です。ただし、配列の現在の長さを超えないように Math.min を使用します。これらの要素をそれぞれ val1, val2 とします。
    • もし val1 < val2 であれば、arr1 の先頭から val1 までの要素(合計 targetK / 2 個)は、K番目に小さい要素にはなりえません。これらの要素は、少なくとも targetK / 2 個の arr2 の要素よりも小さいか等しく、さらに arr1 自身の中にも targetK / 2 - 1 個の要素が存在するため、K番目の要素としては小さすぎるからです。したがって、arr1 の探索範囲を start1 + targetK / 2 からに更新し、targetK の値を targetK - (targetK / 2) に減らして再帰呼び出しします。
    • 逆にもし val1 >= val2 であれば、同様に arr2 の先頭から val2 までの要素(合計 targetK / 2 個)を捨て、arr2 の探索範囲を start2 + targetK / 2 からに更新し、targetK の値を targetK - (targetK / 2) に減らして再帰呼び出しします。

このロジックにより、各ステップで探索範囲が約半分になるため、全体の時間計算量は O(log(m+n)) となります。


class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int lenA = nums1.length;
        int lenB = nums2.length;
        int totalLength = lenA + lenB;

        // 結合された配列の長さが偶数の場合、(totalLength/2) 番目と (totalLength/2 + 1) 番目の要素の平均。
        // 奇数の場合、(totalLength/2 + 1) 番目の要素(両方とも同じK番目を指す)が中央値。
        int kForMedian1 = (totalLength + 1) / 2;
        int kForMedian2 = (totalLength + 2) / 2;

        return (getKthSmallest(nums1, 0, lenA - 1, nums2, 0, lenB - 1, kForMedian1) +
                getKthSmallest(nums1, 0, lenA - 1, nums2, 0, lenB - 1, kForMedian2)) * 0.5;
    }

    /**
     * 二つのソート済み配列からK番目に小さい要素を見つけるヘルパー関数
     * @param arr1 最初の配列
     * @param s1 arr1 の現在の開始インデックス
     * @param e1 arr1 の現在の終了インデックス
     * @param arr2 二番目の配列
     * @param s2 arr2 の現在の開始インデックス
     * @param e2 arr2 の現在の終了インデックス
     * @param targetK 見つけたいK番目の要素 (1-indexed)
     * @return K番目に小さい要素
     */
    private int getKthSmallest(int[] arr1, int s1, int e1,
                                 int[] arr2, int s2, int e2, int targetK) {
        int currentLen1 = e1 - s1 + 1; // arr1 の現在の有効な長さ
        int currentLen2 = e2 - s2 + 1; // arr2 の現在の有効な長さ

        // arr1 が常に短い方の配列になるように交換することで、コードの複雑さを軽減
        if (currentLen1 > currentLen2) {
            return getKthSmallest(arr2, s2, e2, arr1, s1, e1, targetK);
        }

        // arr1 が空の場合、arr2 から targetK 番目の要素を返す
        if (currentLen1 == 0) {
            return arr2[s2 + targetK - 1];
        }

        // targetK が 1 の場合、arr1 と arr2 の現在の先頭要素の小さい方を返す
        if (targetK == 1) {
            return Math.min(arr1[s1], arr2[s2]);
        }

        // arr1 と arr2 からそれぞれ targetK/2 個の要素を比較対象として選ぶ
        // 配列の境界を超えないように Math.min を使用
        int numElementsToCompare1 = Math.min(currentLen1, targetK / 2);
        int numElementsToCompare2 = Math.min(currentLen2, targetK / 2);

        // 比較する要素のインデックス
        int compareIdx1 = s1 + numElementsToCompare1 - 1;
        int compareIdx2 = s2 + numElementsToCompare2 - 1;

        // arr1 の候補要素が arr2 の候補要素より小さい場合
        if (arr1[compareIdx1] < arr2[compareIdx2]) {
            // arr1 の先頭から numElementsToCompare1 個の要素は targetK 番目にはなりえないため捨てる
            // targetK をその分減らして再帰呼び出し
            return getKthSmallest(arr1, compareIdx1 + 1, e1, arr2, s2, e2,
                                   targetK - numElementsToCompare1);
        } else {
            // arr2 の先頭から numElementsToCompare2 個の要素は targetK 番目にはなりえないため捨てる
            // targetK をその分減らして再帰呼び出し
            return getKthSmallest(arr1, s1, e1, arr2, compareIdx2 + 1, e2,
                                   targetK - numElementsToCompare2);
        }
    }
}

タグ: 二分探索 中央値 ソート済み配列 アルゴリズム Java

5月16日 14:19 投稿