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

题目链接

leetcode

题目描述

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。
请你找出并返回这两个正序数组的中位数 。
算法的时间复杂度为 O(log(m + n)) 。

C++ 代码

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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

void printArray(const vector<int>& nums) {
    for (int i = 0; i < nums.size(); i++) {
        cout << nums[i] << " ";
    }
    cout << endl;
}

/*
    不需要合并两个有序数组，只要找到中位数的位置。
    由于两个数组的长度已知，所以中位数对应的两个数组的下标之和也是已知的。
    维护两个指针，初始时分别指向两个数组的下标 0 的位置，
    每次将指向较小值的指针后移一位（若一个指针已经到达数组末尾，则只移动另一个数组的指针），
    直到到达中位数的位置。

    时间复杂度：O(logn)
        其中 m 和 n 分别是数组 nums1 和 nums2 的长度。
        初始时，k = (m+n)/2 或 k = (m+n)/2 + 1，每一轮循环可以将查找范围减少一半，
        因此时间复杂度是 O(log(m+n))。
    空间复杂度：O(1)
*/
class Solution {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int length = nums1.size() + nums2.size();
        if (length % 2 == 0) {
            int left = findKthNumber(nums1, nums2, length / 2);
            int right = findKthNumber(nums1, nums2, length / 2 + 1);
            return (left + right) / 2.0;
        } else {
            return findKthNumber(nums1, nums2, length / 2 + 1);
        }
    }

    int findKthNumber(const vector<int>& nums1, const vector<int>& nums2, int k) {
        int m = nums1.size(), n = nums2.size();
        int index1 = 0, index2 = 0;

        while (true) {
            // 边界情况
            if (index1 == m) {
                return nums2[index2 + k - 1];
            }
            if (index2 == n) {
                return nums1[index1 + k - 1];
            }
            if (k == 1) {
                return min(nums1[index1], nums2[index2]);
            }
            
            // 正常情况
            int newIndex1 = min(index1 + k / 2 - 1, m - 1);
            int newIndex2 = min(index2 + k / 2 - 1, n - 1);
            int pivot1 = nums1[newIndex1];
            int pivot2 = nums2[newIndex2];
            if (pivot1 <= pivot2) {
                k -= newIndex1 - index1 + 1;
                index1 = newIndex1 + 1;
            } else {
                k -= newIndex2 - index2 + 1;
                index2 = newIndex2 + 1;
            }
        }
    }
};

/*
    使用划分的方法解决这个问题，需要理解「中位数的作用是什么」。
    在统计中，中位数被用来：将一个集合划分为两个长度相等的子集，
    其中一个子集中的元素总是大于另一个子集中的元素。

    时间复杂度：O(logn)
        其中 m 和 n 分别是数组 nums1 和 nums2 的长度。
        查找的区间是 [0, m]，该区间的长度在每次循环之后会减少为原来的一半，只需要执行 logm 次循环。
        由于每次循环中的操作次数是常数，所以时间复杂度为 O(logm)。
        由于可能交换 nums1和 nums2 使得 m <= n，所以时间复杂度是 O(log(min(m,n)))。
    空间复杂度：O(1)
*/
class Solution_1 {
public:
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        if (nums1.size() > nums2.size()) {
            return findMedianSortedArrays(nums2, nums1);
        }

        int m = nums1.size(), n = nums2.size();
        int left = 0, right = m;
        // 中位数的左右两个候选值（leftMax：前一部分的最大值，rightMin：后一部分的最小值）
        int leftMax = INT_MIN, rightMin = INT_MAX;
        while (left <= right) {
            // 前一部分包含 nums1[0, i-1] 和 nums2[0, j-1]
            // 后一部分包含 nums1[i, m-1] 和 nums2[j, n-1]
            int i = (left + right) / 2;
            int j = (m + n + 1) / 2 - i;

            int nums_im1 = i == 0 ? INT_MIN : nums1[i - 1];
            int nums_i = i == m ? INT_MAX : nums1[i];
            int nums_jm1 = j == 0 ? INT_MIN : nums2[j - 1];
            int nums_j = j == n ? INT_MAX : nums2[j];
            if (nums_im1 <= nums_j) {
                leftMax = max(nums_im1, nums_jm1);
                rightMin = min(nums_i, nums_j);
                left = i + 1;
            } else {
                right = i - 1;
            }
        }

        return (m + n) % 2 == 0 ? (leftMax + rightMin) / 2.0 : leftMax;
    }
};

int main() {
    Solution solution;
    vector<int> nums1 = {1, 2};
    vector<int> nums2 = {3, 4};
    printArray(nums1);
    printArray(nums2);
    double median = solution.findMedianSortedArrays(nums1, nums2);
    cout << "The median of the two arrays is: " << median << endl;

    return 0;
}

Golang 代码

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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
package main

import (
	"fmt"
	"math"
	"sort"
)

/*
   暴力解法：合并两个数组并排序来找到中位数。
   时间复杂度：O(nlogn)
*/
func findMedianSortedArrays_0(nums1 []int, nums2 []int) float64 {
    // 合并两个数组
	merged := append(nums1, nums2...)
	// 对合并后的数组进行排序
	sort.Ints(merged)

	// 计算中位数
	n := len(merged)
	if n % 2 == 1 {
		// 如果总长度为奇数，中位数是中间的值
		return float64(merged[n/2])
	} else {
		// 如果总长度为偶数，中位数是中间两个值的平均值
		return float64(merged[n/2-1] + merged[n/2]) / 2.0
	}
}

/*
    不需要合并两个有序数组，只要找到中位数的位置。
    由于两个数组的长度已知，所以中位数对应的两个数组的下标之和也是已知的。
    维护两个指针，初始时分别指向两个数组的下标 0 的位置，
    每次将指向较小值的指针后移一位（若一个指针已经到达数组末尾，则只移动另一个数组的指针），
    直到到达中位数的位置。

	时间复杂度：O(logn)
        其中 m 和 n 分别是数组 nums1 和 nums2 的长度。
        初始时，k = (m+n)/2 或 k = (m+n)/2 + 1，每一轮循环可以将查找范围减少一半，
        因此时间复杂度是 O(log(m+n))。
    空间复杂度：O(1)
*/
func findMedianSortedArrays_1(nums1 []int, nums2 []int) float64 {
    m, n := len(nums1), len(nums2)
	totalLength := m + n
	halfLength := totalLength / 2
    
	if totalLength % 2 == 1 {
        // 如果总长度为奇数，中位数是第 halfLength + 1 个元素
		return float64(findKthElement(nums1, nums2, halfLength+1))
	} else {
        // 如果总长度为偶数，中位数是第 halfLength 和 halfLength + 1 个元素的平均值
		return float64(findKthElement(nums1, nums2, halfLength) + findKthElement(nums1, nums2, halfLength+1)) / 2.0
    }
}

// findKthElement 找到两个正序数组中的第 k 小的元素
func findKthElement(nums1, nums2 []int, k int) int {
	i, j := 0, 0
	for {
        // 边界情况
		if i == len(nums1) {
			return nums2[j+k-1]
		}
		if j == len(nums2) {
			return nums1[i+k-1]
		}
        if k == 1 {
			return min(nums1[i], nums2[j])
		}

		// 每次比较两个数组中的第 k/2 个元素
		mid := k / 2
		newI := min(i+mid, len(nums1)) - 1
		newJ := min(j+mid, len(nums2)) - 1

		if nums1[newI] <= nums2[newJ] {
			k -= (newI - i + 1)
			i = newI + 1
		} else {
			k -= (newJ - j + 1)
			j = newJ + 1
		}
	}
}

/*
   二分查找：
       为了满足 O(log(min(m, n))) 的时间复杂度，使用二分查找。
       假设两个数组分别为 A 和 B，长度分别为 m 和 n，且 m <= n。
       在较短的数组 A 中进行二分查找，找到一个分割点 i，
	   使得 A 的左半部分和 B 的左半部分组合后小于等于右半部分。
   分割点的条件：
       分割点 i 和 j 满足：
       i + j = (m + n + 1) / 2（总长度的一半）
       A[i-1] <= B[j] 且 B[j-1] <= A[i]（保证左半部分的最大值小于等于右半部分的最小值）
   计算中位数：
       如果总长度为奇数，中位数是左半部分的最大值。
       如果总长度为偶数，中位数是左半部分的最大值和右半部分的最小值的平均值。

    时间复杂度：O(logn)
        其中 m 和 n 分别是两个数组的长度。
        二分查找：在较短的数组 A 中进行二分查找，时间复杂度为 O(log(min(m,n)))。
    空间复杂度：O(logn)
        常数空间：仅使用了几个变量，空间复杂度为 O(1)。
*/
// findMedianSortedArrays 返回两个正序数组的中位数
func findMedianSortedArrays(nums1 []int, nums2 []int) float64 {
    if len(nums1) > len(nums2) {
        nums1, nums2 = nums2, nums1
    }

    m, n := len(nums1), len(nums2)
    left, right := 0, m
    halfLen := (m + n + 1) / 2

    for left <= right {
        i := (left + right) / 2
        j := halfLen - i

        if i < m && nums1[i] < nums2[j - 1] {
            // 分割点 i 太小，需要增大
            left = i + 1
        } else if i > 0 && nums1[i - 1] > nums2[j] {
            // 分割点 i 太大，需要减小
            right = i - 1
        } else {
            // 找到正确的分割点
            var leftMax float64
            if i == 0 {
                leftMax = float64(nums2[j-1])
            } else if j == 0 {
                leftMax = float64(nums1[i-1])
            } else {
                leftMax = math.Max(float64(nums1[i-1]), float64(nums2[j-1]))
            }

            // 如果总长度是奇数，直接返回 leftMax
            if (m + n) % 2 == 1 {
                return leftMax
            }

            var rightMin float64
            if i == m {
                rightMin = float64(nums2[j])
            } else if j == n {
                rightMin = float64(nums1[i])
            } else {
                rightMin = math.Min(float64(nums1[i]), float64(nums2[j]))
            }

            // 如果总长度是偶数，返回 (leftMax + rightMin) / 2
            return (leftMax + rightMin) / 2
        }
    }
    return 0.0
}

func main() {
	// 测试用例
	testCases := []struct {
		nums1    []int
		nums2    []int
		expected float64
	}{
		{[]int{1, 3}, []int{2}, 2.0},
		{[]int{1, 2}, []int{3, 4}, 2.5},
	}

	for i, tc := range testCases {
		result := findMedianSortedArrays(tc.nums1, tc.nums2)
		fmt.Printf("Test Case %d: Input: nums1 = %v, nums2 = %v\n", i+1, tc.nums1, tc.nums2)
		
        if result == tc.expected {
			fmt.Printf("Output: %.1f, PASS\n", result)
		} else {
			fmt.Printf("Output: %.1f, FAIL (Expected: %.1f)\n", result, tc.expected)
		}
	}
}
#leetcode
leetcode4：寻找两个正序数组的中位数
https://lcf163.github.io/2023/08/09/leetcode4：寻找两个正序数组的中位数/
作者
乘风的小站
发布于
2023年8月9日
许可协议