查找两个已排序数组的中位数(Median of Two Sorted Arrays)

Posted 2020-06-08

https://leetcode.com/problems/median-of-two-sorted-arrays/

一个常见的思路就是利用递归查找第K个数来实现。

 1         private int FindKth(int[] a, int sa, int[] b, int sb, int k)
 2         {
 3             if (a.Length - sa < b.Length - sb)
 4             {
 5                 return FindKth(b, sb, a, sa, k);
 6             }
 7 
 8             if (sb == b.Length)
 9             {
10                 return a[k - 1];
11             }
12 
13             if (k == 1)
14             {
15                 return Math.Min(a[sa], b[sb]);
16             }
17 
18             int j = Math.Min(b.Length - sb, k / 2);
19             int i = k - j;
20 
21             if (b[sb + j - 1] < a[sa + i - 1])
22             {
23                 return FindKth(a, sa, b, sb + j, k - j);
24             }
25             else
26             {
27                 return FindKth(a, sa + i, b, sb, j);
28             }
29         }

 

分析：假设要找两个数组中第k大的数，我们可以假设前k/2个数存在于数组1，剩下的k - k/2个数存在于数组2，也就是18-19行：

            int j = Math.Min(b.Length - sb, k / 2);
            int i = k - j;

然后比较两个数组这两部分的最后一个元素a[sa + i - 1]和b[sb + j -1]，也就是21 ~28行:

            if (b[sb + j - 1] < a[sa + i - 1])
            {
                return FindKth(a, sa, b, sb + j, k - j);
            }
            else
            {
                return FindKth(a, sa + i, b, sb, j);
            }

1. 如果b的最后一个元素较小，那么说明这j个数肯定存在于最小的前k个数里，且第k小的数肯定不在这j个数里，否则b的最后一个元素应该不仅比前j-1个元素大，而且应该比a中的前k - j个元素大，这样它才是两个数组中第k大的元素，但b[sb + j - 1] < a[sa + i - 1]，矛盾。

这样我们只需继续在两个数组剩下的元素里找到最小的k - j个元素，我们就找到了第k小的元素。

2. 反之，如果a[sa + i - 1] <= b[sb + j -1], 说明a数组里的部分是两个数组两部分里较小的k - j个元素，第k个元素也不存在于这个部分，所以我们只需在两数组剩余的部分查找第j大的元素即可。

 

接下来解释一下开头的几个判断。

            if (a.Length - sa < b.Length - sb)
            {
                return FindKth(b, sb, a, sa, k);
            }

我们始终把剩余长度较大的数组作为第一个数组，这样是为了简化后续判断逻辑。

            if (sb == b.Length)
            {
                return a[k - 1];
            }

在前一个递归调用中，b部分的长度已经不足k/2，所以我们只需在这次调用中直接返回a的第k-j大元素 (k - j也就是本次调用的参数k，所以是a[k-1])，而无须综合考虑两个数组。

            if (k == 1)
            {
                return Math.Min(a[sa], b[sb]);
            }

循环至最后一个元素时，返回a，b剩余部分中较小的元素作为递归终止。

完整代码：

 1 using System;
 2 
 3 namespace ProblemSetConsoleApp
 4 {
 5     public class LC004_Median_of_Two_Sorted_Arrays
 6     {
 7         public double FindMedianSortedArrays(int[] nums1, int[] nums2)
 8         {
 9             if (nums1 == null || nums1.Length == 0)
10             {
11                 return GetMedianFromArray(nums2);
12             }
13 
14             if (nums2 == null || nums2.Length == 0)
15             {
16                 return GetMedianFromArray(nums1);
17             }
18 
19             int totalLen = nums1.Length + nums2.Length;
20 
21             if (totalLen == 2)
22             {
23                 return (nums1[0] + nums2[0]) / 2.0;
24             }
25 
26             double m1 = FindKth(nums1, 0, nums2, 0, (totalLen + 1) / 2);
27 
28             if ((totalLen & 1) == 0)
29             {
30                 int m2 = FindKth(nums1, 0, nums2, 0, totalLen / 2 + 1);
31                 m1 = (m1 + m2) / 2.0;
32             }
33 
34             return m1;
35         }
36 
37         private int FindKth(int[] a, int sa, int[] b, int sb, int k)
38         {
39             if (a.Length - sa < b.Length - sb)
40             {
41                 return FindKth(b, sb, a, sa, k);
42             }
43 
44             if (sb == b.Length)
45             {
46                 return a[k - 1];
47             }
48 
49             if (k == 1)
50             {
51                 return Math.Min(a[sa], b[sb]);
52             }
53 
54             int j = Math.Min(b.Length - sb, k / 2);
55             int i = k - j;
56 
57             if (b[sb + j - 1] < a[sa + i - 1])
58             {
59                 return FindKth(a, sa, b, sb + j, k - j);
60             }
61             else
62             {
63                 return FindKth(a, sa + i, b, sb, j);
64             }
65         }
66 
67         private double GetMedianFromArray(int[] a)
68         {
69             if (a != null && a.Length > 0)
70             {
71                 int m = a.Length / 2;
72                 if ((a.Length & 1) != 0)
73                 {
74                     return a[m];
75                 }
76                 else
77                 {
78                     return (a[m] + a[m - 1]) / 2.0;
79                 }
80             }
81 
82             return 0.0;
83         }
84     }
85 }

View Code

测试用例：

 1 using ProblemSetConsoleApp;
 2 using Microsoft.VisualStudio.TestTools.UnitTesting;
 3 
 4 namespace ProblemSetConsoleAppTest
 5 {
 6     /// <summary>
 7     /// Summary description for LC004_Median_of_Two_Sorted_ArraysTest
 8     /// </summary>
 9     [TestClass]
10     public class LC004_Median_of_Two_Sorted_ArraysTest
11     {
12         static LC004_Median_of_Two_Sorted_Arrays seeker = new LC004_Median_of_Two_Sorted_Arrays();
13 
14         [TestMethod]
15         public void TestOddOdd()
16         {
17             int[] a = { 1, 3, 5, 7, 9 };
18             int[] b = { 2, 4, 6, 8, 10 };
19             var m = seeker.FindMedianSortedArrays(a, b);
20             Assert.AreEqual(5.5, m);
21         }
22 
23         [TestMethod]
24         public void TestOddEven()
25         {
26             int[] a = { 1, 3, 5, 7, 9 };
27             int[] b = { 2, 4, 6, 8 };
28             Assert.AreEqual(5, seeker.FindMedianSortedArrays(a, b));
29         }
30 
31         [TestMethod]
32         public void TestEvenEven()
33         {
34             int[] a = { 1, 3, 5, 7 };
35             int[] b = { 2, 4, 6, 8 };
36             Assert.AreEqual(4.5, seeker.FindMedianSortedArrays(a, b));
37         }
38 
39         [TestMethod]
40         public void TestSingleElement()
41         {
42             int[] a = { 1 };
43             int[] b = { 2 };
44             Assert.AreEqual(1.5, seeker.FindMedianSortedArrays(a, b));
45         }
46 
47         [TestMethod]
48         public void TestSingleElement2()
49         {
50             int[] a = { 1 };
51             int[] b = { 2, 3 };
52             Assert.AreEqual(2, seeker.FindMedianSortedArrays(a, b));
53         }
54 
55         [TestMethod]
56         public void TestSingleElement3()
57         {
58             int[] a = { 1, 2 };
59             int[] b = { 3 };
60             Assert.AreEqual(2, seeker.FindMedianSortedArrays(a, b));
61         }
62 
63         [TestMethod]
64         public void TestSmall1()
65         {
66             int[] a = { 1, 2 };
67             int[] b = { 3, 4 };
68             Assert.AreEqual(2.5, seeker.FindMedianSortedArrays(a, b));
69         }
70 
71         [TestMethod]
72         public void TestSmall2()
73         {
74             int[] a = { 3, 4 };
75             int[] b = { 1, 2 };
76             Assert.AreEqual(2.5, seeker.FindMedianSortedArrays(a, b));
77         }
78     }
79 }

View Code