寻找两个正序数组的中位数

题意

给定两个正序数组，找中位数。要求 $O(\log(m+n))$ 。

思路

中位数等价于把合并后数组分成左右两半，使得左半最大 $\leqslant$ 右半最小，且左半元素数 $=$ 右半元素数（或恰好多一个）。

在较短数组 a 上二分分割线位置 i（a 左半取 a[0..i-1]），b 的分割位置 j = (m+n+1)/2 - i 自动确定。四个边界值 al, ar, bl, br 分别表示分割线两侧的值（越界用 $\pm\infty$ ）。

条件 al <= br && bl <= ar 满足时分割合法，中位数由 max(al,bl) 和 min(ar,br) 计算。al > br 时 i 太大，bl > ar 时 i 太小。

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    double findMedianSortedArrays(vector<int> &a, vector<int> &b) {
        if (a.size() > b.size())
            swap(a, b);
        int m = a.size(), n = b.size(), lo = 0, hi = m;
        while (lo <= hi) {
            int i = (lo + hi) / 2, j = (m + n + 1) / 2 - i;
            int al = (i == 0 ? INT_MIN : a[i - 1]), ar = (i == m ? INT_MAX : a[i]);
            int bl = (j == 0 ? INT_MIN : b[j - 1]), br = (j == n ? INT_MAX : b[j]);
            if (al <= br && bl <= ar) {
                if ((m + n) % 2)
                    return max(al, bl);
                return (max(al, bl) + min(ar, br)) / 2.0;
            }
            if (al > br)
                hi = i - 1;
            else
                lo = i + 1;
        }
        return 0;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int m, n;
    cin >> m >> n;
    vector<int> a(m), b(n);
    for (int &x : a)
        cin >> x;
    for (int &x : b)
        cin >> x;
    cout << fixed << setprecision(5) << Solution().findMedianSortedArrays(a, b) << '\n';
    return 0;
}

python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35#!/usr/bin/env python3
from typing import List


class Solution:
    def findMedianSortedArrays(self, a: List[int], b: List[int]) -> float:
        if len(a) > len(b):
            a, b = b, a
        m, n = len(a), len(b)
        lo, hi = 0, m
        while lo <= hi:
            i = (lo + hi) // 2
            j = (m + n + 1) // 2 - i
            al = a[i - 1] if i > 0 else float("-inf")
            ar = a[i] if i < m else float("inf")
            bl = b[j - 1] if j > 0 else float("-inf")
            br = b[j] if j < n else float("inf")
            if al <= br and bl <= ar:
                return max(al, bl) if (m + n) % 2 else (max(al, bl) + min(ar, br)) / 2
            if al > br:
                hi = i - 1
            else:
                lo = i + 1
        return 0


def main():
    m, n = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))
    print(Solution().findMedianSortedArrays(a, b))


if __name__ == "__main__":
    main()

复杂度

时间复杂度： $O(\log(\min(m,n)))$ ，只在较短数组上二分。
空间复杂度： $O(1)$ 。

总结

双数组中位数的核心是"分割线"模型：在较短数组上二分，自动确定另一个数组的分割位置，四个边界值判断合法性并计算结果。j = (m+n+1)/2 - i 保证了左半元素数正确。

寻找两个正序数组的中位数

目录

题意

思路

代码

复杂度

总结