LeetCode 兩個有序數組的中位數 - 二分查找第 k 小

問題描述

給定兩個已排序數組 A、B，長度分別為 m、n。求合併後的中位數。要求時間複雜度 O(log(m+n))。

轉化為求第 k 小元素

• 若總長度 L = m+n 為奇數，中位數為第 (L+1)/2 小。
• 若為偶數，中位數為第 L/2 小與第 L/2+1 小的平均。

二分查找第 k 小

比較 A[k/2-1] 和 B[k/2-1]（邊界處理）：

• 若 A[k/2-1] < B[k/2-1]，則 A 前 k/2 個元素不可能是第 k 小，丟棄，遞歸在 A[k/2..] 和 B 中找第 k-k/2 小。
• 反之丟棄 B 前 k/2 個。
• 相等則直接返回該值。

基準情況：一數組為空直接返回另一數組第 k-1 個；k=1 返回兩數組首元素較小者。

代碼實現（C++）

class Solution {
public:
    double findMedianSortedArrays(vector& A, vector& B) {
        int m = A.size(), n = B.size();
        int total = m + n;
        if (total % 2 == 1) return findKth(A, 0, B, 0, (total+1)/2);
        else return (findKth(A, 0, B, 0, total/2) + findKth(A, 0, B, 0, total/2+1)) / 2.0;
    }
    double findKth(vector& A, int i, vector& B, int j, int k) {
        if (i >= A.size()) return B[j + k - 1];
        if (j >= B.size()) return A[i + k - 1];
        if (k == 1) return min(A[i], B[j]);
        int step = k / 2;
        int aVal = (i + step - 1 < A.size()) ? A[i + step - 1] : INT_MAX;
        int bVal = (j + step - 1 < B.size()) ? B[j + step - 1] : INT_MAX;
        if (aVal < bVal) return findKth(A, i + step, B, j, k - step);
        else return findKth(A, i, B, j + step, k - step);
    }
};

關鍵點

• 使用索引偏移而非數組切片，避免拷貝開銷。
• INT_MAX 模擬越界為無窮大，保證另一數組被選中。
• 遞歸深度 O(log k)，總時間 O(log(m+n))。