問題概要
二つの昇順にソートされた整数配列 nums1 と nums2 が与えられます。それぞれの配列のサイズは m と n です。これら二つの配列を結合した場合の中央値を求めてください。
このアルゴリズムの時間計算量は O(log (m+n)) である必要があります。
例1:
入力:nums1 = [1,3], nums2 = [2]
出力:2.00000
解説:結合配列 = [1,2,3] 、中央値 2
アプローチ1:マージによる方法 (時間計算量 O(m+n))
この問題に対する最も直接的なアプローチは、二つの与えられたソート済み配列 nums1 と nums2 を一つの大きなソート済み配列にマージすることです。マージが完了したら、その結合された配列から中央値を計算します。
- マージ処理は、両配列の先頭から要素を順に比較し、小さい方を新しい配列に追加していくことで行われます。
- 結合された配列の長さが奇数の場合、中央値は単純に配列の中央要素です。
- 結合された配列の長さが偶数の場合、中央値は中央の二つの要素の平均となります。
- この方法の時間計算量は
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番目に小さい要素を見つける再帰的ロジック
二つの配列 arr1、arr2 と、それぞれの探索範囲(start1, end1, start2, end2)、そして目標とする順序 targetK を引数に取る再帰関数 getKthSmallest を設計します。
基本ケース:
- もし
arr1の探索範囲が空であれば(currentLen1 == 0)、arr2の中からtargetK番目に小さい要素(start2 + targetK - 1インデックス)を返します。arr2が空の場合も同様です。 - もし
targetKが1であれば、arr1[start1]とarr2[start2]のうち小さい方が1番目に小さい要素なので、それを返します。 - 再帰呼び出しを安定させるため、
arr1が常に短い方の配列になるように、必要に応じて配列を入れ替えて呼び出します。
- もし
探索範囲の縮小:
- それぞれの配列から
targetK / 2個の要素を選択し、それらを中央値の候補として比較します。具体的には、arr1のstart1 + targetK / 2 - 1番目の要素と、arr2のstart2 + 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);
}
}
}