4. 两个正序数组的中位数

题目描述

输入两个已按从小到大排好序的数组 nums1 和 nums2，长度分别为 m 和 n。

把两个数组的全部元素放在一起看，返回它们的中位数。元素总数为奇数时，中位数是排序后正中间那个数；总数为偶数时，中位数是正中间两个数的平均值。

要求算法的时间复杂度为 O(log (m+n))。

示例 1

输入：nums1 = [1,3], nums2 = [2]
输出：2.0

解释：合并后为 [1,2,3]，正中间是 2。

示例 2

输入：nums1 = [1,2], nums2 = [3,4]
输出：2.5

解释：合并后为 [1,2,3,4]，中间两个数是 2 和 3，平均值为 2.5。

数据范围

nums1.length == m
nums2.length == n
0 <= m <= 1000
0 <= n <= 1000
1 <= m + n <= 2000
-10^6 <= nums1[i], nums2[i] <= 10^6

解题思路：二分划分较短数组

在较短数组上二分切分。

为什么正确

切分两侧交叉边界有序。

复杂度分析

时间 O(log min(m,n))，额外空间 O(1)。

易错点

注意边界条件并在回溯或弹出后恢复状态。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
 def findMedianSortedArrays(self,A,B):
  if len(A)>len(B):A,B=B,A
  m,n=len(A),len(B);l,h=0,m
  while l<=h:
   i=(l+h)//2;j=(m+n+1)//2-i;al=A[i-1] if i else float('-inf');ar=A[i] if i<m else float('inf');bl=B[j-1] if j else float('-inf');br=B[j] if j<n else float('inf')
   if al<=br and bl<=ar:return float(max(al,bl)) if (m+n)%2 else (max(al,bl)+min(ar,br))/2
   if al>br:h=i-1
   else:l=i+1

class Solution {
    public double findMedianSortedArrays(int[] A, int[] B) {
        if (A.length > B.length) { int[] t = A; A = B; B = t; }
        int m = A.length, n = B.length, l = 0, h = m;
        while (l <= h) {
            int i = (l + h) / 2, j = (m + n + 1) / 2 - i;
            double al = i > 0 ? A[i - 1] : Double.NEGATIVE_INFINITY, ar = i < m ? A[i] : Double.POSITIVE_INFINITY;
            double bl = j > 0 ? B[j - 1] : Double.NEGATIVE_INFINITY, br = j < n ? B[j] : Double.POSITIVE_INFINITY;
            if (al <= br && bl <= ar) return (m + n) % 2 == 1 ? Math.max(al, bl) : (Math.max(al, bl) + Math.min(ar, br)) / 2;
            if (al > br) h = i - 1;
            else l = i + 1;
        }
        return 0;
    }
}

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(), l = 0, h = m;
        while (l <= h) {
            int i = (l + h) / 2, j = (m + n + 1) / 2 - i;
            double al = i ? A[i - 1] : -INFINITY, ar = i < m ? A[i] : INFINITY;
            double bl = j ? B[j - 1] : -INFINITY, br = j < n ? B[j] : INFINITY;
            if (al <= br && bl <= ar) return (m + n) % 2 ? max(al, bl) : (max(al, bl) + min(ar, br)) / 2;
            if (al > br) h = i - 1;
            else l = i + 1;
        }
        return 0;
    }
};

func findMedianSortedArrays(A []int, B []int) float64 {
    if len(A) > len(B) {
        A, B = B, A
    }
    m, n := len(A), len(B)
    l, h := 0, m
    for l <= h {
        i := (l + h) / 2
        j := (m+n+1)/2 - i
        al, ar, bl, br := math.Inf(-1), math.Inf(1), math.Inf(-1), math.Inf(1)
        if i > 0 { al = float64(A[i-1]) }
        if i < m { ar = float64(A[i]) }
        if j > 0 { bl = float64(B[j-1]) }
        if j < n { br = float64(B[j]) }
        if al <= br && bl <= ar {
            if (m+n)%2 == 1 { return math.Max(al, bl) }
            return (math.Max(al, bl) + math.Min(ar, br)) / 2
        }
        if al > br {
            h = i - 1
        } else {
            l = i + 1
        }
    }
    return 0
}

var findMedianSortedArrays = function(A, B) {
    if (A.length > B.length) [A, B] = [B, A];
    const m = A.length, n = B.length; let l = 0, h = m;
    while (l <= h) {
        const i = (l + h) >> 1, j = ((m + n + 1) >> 1) - i;
        const al = i ? A[i - 1] : -Infinity, ar = i < m ? A[i] : Infinity;
        const bl = j ? B[j - 1] : -Infinity, br = j < n ? B[j] : Infinity;
        if (al <= br && bl <= ar) return (m + n) % 2 ? Math.max(al, bl) : (Math.max(al, bl) + Math.min(ar, br)) / 2;
        if (al > br) h = i - 1;
        else l = i + 1;
    }
    return 0;
};

function findMedianSortedArrays(A: number[], B: number[]): number {
    if (A.length > B.length) [A, B] = [B, A];
    const m = A.length, n = B.length; let l = 0, h = m;
    while (l <= h) {
        const i = (l + h) >> 1, j = ((m + n + 1) >> 1) - i;
        const al = i ? A[i - 1] : -Infinity, ar = i < m ? A[i] : Infinity;
        const bl = j ? B[j - 1] : -Infinity, br = j < n ? B[j] : Infinity;
        if (al <= br && bl <= ar) return (m + n) % 2 ? Math.max(al, bl) : (Math.max(al, bl) + Math.min(ar, br)) / 2;
        if (al > br) h = i - 1;
        else l = i + 1;
    }
    return 0;
};

载入编辑器，自己跑一遍 →

其他解法

归并法 O(m+n) 最直白；第 k 小二分每轮排除一半候选，达到 O(log(m+n))；主解法只在较短数组上二分切分位置，达到 O(log min(m,n))。

解法 2：二分找第 k 小时间 O(log(m+n))，空间 O(1)

把中位数转化为求两个数组合并后的第 k 小元素。每轮比较 A 和 B 各自第 k/2 个候选元素，较小的一方连同它之前的元素都不可能是第 k 小，整段排除并减小 k。某个数组取完或 k = 1 时直接得到答案。

思路比划分更容易想到，边界也较少。总长度为偶数时要求两次第 k 小，复杂度是 O(log(m+n))，略逊于主解法的 O(log min(m,n))。

class Solution:
    def findMedianSortedArrays(self, A, B):
        m, n = len(A), len(B)

        def kth(k):
            i = j = 0
            while True:
                if i == m:
                    return B[j + k - 1]
                if j == n:
                    return A[i + k - 1]
                if k == 1:
                    return min(A[i], B[j])
                half = k // 2
                ni = min(i + half, m) - 1
                nj = min(j + half, n) - 1
                if A[ni] <= B[nj]:
                    k -= ni - i + 1
                    i = ni + 1
                else:
                    k -= nj - j + 1
                    j = nj + 1

        total = m + n
        if total % 2:
            return float(kth(total // 2 + 1))
        return (kth(total // 2) + kth(total // 2 + 1)) / 2

class Solution {
    public double findMedianSortedArrays(int[] A, int[] B) {
        int total = A.length + B.length;
        if (total % 2 == 1) return kth(A, B, total / 2 + 1);
        return (kth(A, B, total / 2) + kth(A, B, total / 2 + 1)) / 2.0;
    }

    private int kth(int[] A, int[] B, int k) {
        int m = A.length, n = B.length, i = 0, j = 0;
        while (true) {
            if (i == m) return B[j + k - 1];
            if (j == n) return A[i + k - 1];
            if (k == 1) return Math.min(A[i], B[j]);
            int half = k / 2;
            int ni = Math.min(i + half, m) - 1;
            int nj = Math.min(j + half, n) - 1;
            if (A[ni] <= B[nj]) {
                k -= ni - i + 1;
                i = ni + 1;
            } else {
                k -= nj - j + 1;
                j = nj + 1;
            }
        }
    }
}

class Solution {
public:
    double findMedianSortedArrays(vector<int>& A, vector<int>& B) {
        int total = A.size() + B.size();
        if (total % 2) return kth(A, B, total / 2 + 1);
        return (kth(A, B, total / 2) + kth(A, B, total / 2 + 1)) / 2.0;
    }

    int kth(vector<int>& A, vector<int>& B, int k) {
        int m = A.size(), n = B.size(), i = 0, j = 0;
        while (true) {
            if (i == m) return B[j + k - 1];
            if (j == n) return A[i + k - 1];
            if (k == 1) return min(A[i], B[j]);
            int half = k / 2;
            int ni = min(i + half, m) - 1;
            int nj = min(j + half, n) - 1;
            if (A[ni] <= B[nj]) {
                k -= ni - i + 1;
                i = ni + 1;
            } else {
                k -= nj - j + 1;
                j = nj + 1;
            }
        }
    }
};

func findMedianSortedArrays(A []int, B []int) float64 {
	total := len(A) + len(B)
	if total%2 == 1 {
		return float64(kth(A, B, total/2+1))
	}
	return float64(kth(A, B, total/2)+kth(A, B, total/2+1)) / 2.0
}

func kth(A, B []int, k int) int {
	m, n, i, j := len(A), len(B), 0, 0
	for {
		if i == m {
			return B[j+k-1]
		}
		if j == n {
			return A[i+k-1]
		}
		if k == 1 {
			return min(A[i], B[j])
		}
		half := k / 2
		ni := min(i+half, m) - 1
		nj := min(j+half, n) - 1
		if A[ni] <= B[nj] {
			k -= ni - i + 1
			i = ni + 1
		} else {
			k -= nj - j + 1
			j = nj + 1
		}
	}
}

var findMedianSortedArrays = function(A, B) {
    const m = A.length, n = B.length;
    const kth = (k) => {
        let i = 0, j = 0;
        while (true) {
            if (i === m) return B[j + k - 1];
            if (j === n) return A[i + k - 1];
            if (k === 1) return Math.min(A[i], B[j]);
            const half = Math.floor(k / 2);
            const ni = Math.min(i + half, m) - 1;
            const nj = Math.min(j + half, n) - 1;
            if (A[ni] <= B[nj]) {
                k -= ni - i + 1;
                i = ni + 1;
            } else {
                k -= nj - j + 1;
                j = nj + 1;
            }
        }
    };
    const total = m + n;
    if (total % 2 === 1) return kth(Math.floor(total / 2) + 1);
    return (kth(total / 2) + kth(total / 2 + 1)) / 2;
};

function findMedianSortedArrays(A: number[], B: number[]): number {
    const m = A.length, n = B.length;
    const kth = (k: number): number => {
        let i = 0, j = 0;
        while (true) {
            if (i === m) return B[j + k - 1];
            if (j === n) return A[i + k - 1];
            if (k === 1) return Math.min(A[i], B[j]);
            const half = Math.floor(k / 2);
            const ni = Math.min(i + half, m) - 1;
            const nj = Math.min(j + half, n) - 1;
            if (A[ni] <= B[nj]) {
                k -= ni - i + 1;
                i = ni + 1;
            } else {
                k -= nj - j + 1;
                j = nj + 1;
            }
        }
    };
    const total = m + n;
    if (total % 2 === 1) return kth(Math.floor(total / 2) + 1);
    return (kth(total / 2) + kth(total / 2 + 1)) / 2;
};

解法 3：双指针归并到中点时间 O(m+n)，空间 O(1)

用两个指针按归并顺序从小到大取元素，只走到第 (m+n)/2 个位置。过程中保存最后两个取出的值 prev 和 cur。总长度为奇数时答案是 cur，为偶数时是 (prev + cur) / 2。

实现最简单，不容易出错。但时间是线性的，不满足题目对 O(log(m+n)) 的要求，适合作为对拍基准。

class Solution:
    def findMedianSortedArrays(self, A, B):
        m, n = len(A), len(B)
        total = m + n
        i = j = 0
        prev = cur = 0
        for _ in range(total // 2 + 1):
            prev = cur
            if i < m and (j >= n or A[i] <= B[j]):
                cur = A[i]
                i += 1
            else:
                cur = B[j]
                j += 1
        if total % 2:
            return float(cur)
        return (prev + cur) / 2

class Solution {
    public double findMedianSortedArrays(int[] A, int[] B) {
        int m = A.length, n = B.length, total = m + n;
        int i = 0, j = 0, prev = 0, cur = 0;
        for (int step = 0; step <= total / 2; step++) {
            prev = cur;
            if (i < m && (j >= n || A[i] <= B[j])) cur = A[i++];
            else cur = B[j++];
        }
        if (total % 2 == 1) return cur;
        return (prev + cur) / 2.0;
    }
}

class Solution {
public:
    double findMedianSortedArrays(vector<int>& A, vector<int>& B) {
        int m = A.size(), n = B.size(), total = m + n;
        int i = 0, j = 0, prev = 0, cur = 0;
        for (int step = 0; step <= total / 2; step++) {
            prev = cur;
            if (i < m && (j >= n || A[i] <= B[j])) cur = A[i++];
            else cur = B[j++];
        }
        if (total % 2) return cur;
        return (prev + cur) / 2.0;
    }
};

func findMedianSortedArrays(A []int, B []int) float64 {
	m, n := len(A), len(B)
	total := m + n
	i, j, prev, cur := 0, 0, 0, 0
	for step := 0; step <= total/2; step++ {
		prev = cur
		if i < m && (j >= n || A[i] <= B[j]) {
			cur = A[i]
			i++
		} else {
			cur = B[j]
			j++
		}
	}
	if total%2 == 1 {
		return float64(cur)
	}
	return float64(prev+cur) / 2.0
}

var findMedianSortedArrays = function(A, B) {
    const m = A.length, n = B.length, total = m + n;
    let i = 0, j = 0, prev = 0, cur = 0;
    for (let step = 0; step <= Math.floor(total / 2); step++) {
        prev = cur;
        if (i < m && (j >= n || A[i] <= B[j])) cur = A[i++];
        else cur = B[j++];
    }
    if (total % 2 === 1) return cur;
    return (prev + cur) / 2;
};

function findMedianSortedArrays(A: number[], B: number[]): number {
    const m = A.length, n = B.length, total = m + n;
    let i = 0, j = 0, prev = 0, cur = 0;
    for (let step = 0; step <= Math.floor(total / 2); step++) {
        prev = cur;
        if (i < m && (j >= n || A[i] <= B[j])) cur = A[i++];
        else cur = B[j++];
    }
    if (total % 2 === 1) return cur;
    return (prev + cur) / 2;
};

二分查找专题的其他题目

查看二分查找全部题目 →