25 Mar 2018 • on 计算机 C++

求两个有序数列中位数算法

leetcode 上第四题, 两个有序数列, 长度分别为 m 和 n, 求所有元素中的中位数.看上去很简单, 排序就行, 不过时间复杂度是(m + n)log(m + n). 这里解释一下log(min(m, n))复杂度的算法.

什么是中位数

对于求中位数, 想到的自然是奇数个元素取中间的元素, 偶数个元素就取中间两个元素除以2. 再深入一下, 所谓的中位数相当于把一个有序数列(下面均看作是递增数列)分为均等的两部分: 左面的数都比中位数小, 右面的数都比中位数大. 算法的核心就是根据这个定义来进行的.

划分两个数列

根据上面的定义, 需要把两个有序数列分为均等的两部分, 分割完应该是下面这样:

A[0] ... A[i - 1] | A[i] ... A[m - 1]
B[0] ... B[j - 1] | B[j] ... B[n - 1]

即以 A[i] 和 B[j] 为边界划分, 这里先不考虑边界是否存在的问题, 最后会统一考虑. 那么当(m + n)是奇数时, 我们可以让左面比右面多一个元素, 则中位数为左面元素的最大值;(m + n)是偶数时, 中位数就应该是左面的最大值和右面最小值之和的一半:

if (m + n) % 2 == 1 {
    median = max(A[i - 1], B[j - 1])
} else {} 
    median = (max(A[i - 1], B[j - 1]) + min(A[i], B[j])) / 2
}

i 和 j 如何确定

要使左面任意数字小于右面任意数字, 只需让左面的最大值小于右面的最小值, 因为是有序数列, 所以只需满足:

A[i - 1] < B[j] && B[j - 1] < A[i]

把 i 和 j 化为一个变量

j 是可以用 i 来表示的, 上面分割图可以看出:

i + j = (m + n) / 2
j = (m + n) / 2 - i

从这个式子可以看出, 当m > n时, 如果 i 取值过大, j 可能会出现负值, 所以当m > n时, 应该调换两数列的位置.

为了满足上面所说的, 当(m + n)为奇数时, 左面比右面多一个元素, 可以利用整型除法的舍入, 将上式变成这样:

j = (m + n + 1) / 2 - i

如何求 i

现在问题里只剩一个未知数 i 了, 问题的描述就变为: 在 0 到 m 范围(i 可以取到 m, 即数列A都在左面)求 i, 使 i 满足下面条件:

A[i - 1] < B[(m + n + 1) / 2 - i] && B[(m + n + 1) / 2 - i - 1] < A[i]

使用二分查找的话, 时间复杂度为logm, 因为m < n, 所以该算法时间复杂度为log(min(m, n)).

如何二分查找

上面确定了 i 的取值 iMin == 0, iMax == m. 那么开始在中间取 i 值, i = (iMin + iMax) / 2. 可能出现的情况如下:

A[i] < B[j - 1]

说明 i 太大了, 左面有值大于右面, 应该使 i 左移, 即iMax = i--
A[i - 1] > B[j]

同理, i 太小了, 令iMin = i++

边界处理

最后就是向来头疼的边界问题了:

i == 0

说明数列A都在右面, 所以左面最大值为B[j - 1]
j == 0

同理, 数列B都在右面, 左面最大值为A[i - 1]
i == m

说明数列A都在左面, 右面最小值为B[j]
j == n

同理, 数列B都在左面, 右面最小值为A[i]

最后贴上 C++ 的代码

不得不吐槽一下 leetcode 的提交, 考虑了半天同样的算法, 为什么 dalao 的代码跑起来比我快了三倍, 结果是人家优化了一下 cin……

double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
    int length1 = nums1.size();
    int length2 = nums2.size();
    if(length1 > length2) {
        return findMedianSortedArrays(nums2, nums1);
    }
    int iMin = 0;
    int iMax = length1;
    int halfOfNums = (length1 + length2 + 1) / 2;
    int i, j;
    while(iMin <= iMax) {
        i = (iMin + iMax) / 2;
        j = halfOfNums - i;
        if(i < length1 && nums1[i] < nums2[j - 1]) {
            iMin = i + 1;
        } else if(i > 0 && nums1[i - 1] > nums2[j]) {
            iMax = i - 1;
        } else {
            break;
        }
    }
    int maxOfLeft, minOfRight;
    if(i == 0) {
        maxOfLeft = nums2[j - 1];
    } else if(j == 0) {
        maxOfLeft = nums1[i - 1];
    } else {
        maxOfLeft = max(nums1[i - 1], nums2[j - 1]);
    }
    if((length1 + length2) % 2 == 1) {
        return maxOfLeft;
    }
    if(i == length1) {
        minOfRight = nums2[j];
    } else if(j == length2) {
        minOfRight = nums1[i];
    } else {
        minOfRight = min(nums1[i], nums2[j]);
    }
    return (maxOfLeft + minOfRight) / 2.0;
}

shino