[数据结构与算法]力扣第4题-Swift题解：寻找两个正序数组的中位数

双数组二分查找

题目描述

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的 中位数 。

示例 1：

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

输出： 2.00000

解释： 合并数组 = [1,2,3] ，中位数 2

示例2：

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

输出： 2.50000

解释： 合并数组 = [1,2,3,4] ，中位数 (2 + 3) / 2 = 2.5

示例3：

输入： nums1 = [0,0], nums2 = [0,0]

输出： 0.00000

示例 4：

输入： nums1 = [], nums2 = [1]

输出： 1.00000

示例5：

输入： nums1 = [2], nums2 = []

输出： 2.00000

提示：

• nums1.length == m
• nums2.length == n
• 0 <= m <= 1000
• 0 <= n <= 1000
• 1 <= m + n <= 2000
• -106 <= nums1[i], nums2[i] <= 106

来源：力扣（LeetCode）

链接：https://leetcode-cn.com/problems/median-of-two-sorted-arrays

著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

分析

该问题如果不对时间复杂度和空间复杂度提出要求的话，通过最直接的方法可以通过。
考虑到问题规模，m + n <= 2000，我们设 N = m + n ，这个数据规模非常小。所以即使是 O(n*log(n)) 时间复杂度也可以顺利通过。

该问题按照时间复杂度，空间复杂度的优化程度分别有以下几种解法。

解法 1： 直接合并排序 + 查找

这是最简单的思路。步骤非常的明确。

1. 将数组nums1和nums2合并为新的数组nums3。
2. 用系统函数为nums3进行排序。
3. 寻找nums3中位数。

nums1和nums2合并的时间复杂度为O(N)，排序复杂度为O(N*log(N))，由于需要另外开辟一个空间进行合并和排序，所以空间复杂度为O(N)。
最终在已经排好序的nums3上查找中位数，只要简单的进行中位索引取值即可，这部分复杂度为O(1)。

综上，该算法的总时间复杂度为O(N*log(N))，空间复杂度为O(N)。

解法2：归并合并排序 + 查找

解法2和解法1思路一样，也是合并nums1和nums2形成新的排序数组nums3，区别在于合并的过程。

由于已知nums1和nums2是各自已排序的数组，通过经典排序算法 归并排序 中的合并步骤就可以它们合并成新的排序数组。而该合并步骤的时间复杂度是O(N)。

剩余步骤和解法1相同。

综上，该算法的总时间复杂度为O(N)，空间复杂度为O(N)。

解法 3：双数组二分查找 方法 1

已知数组nums1和nums2都是已排序数组，自然不难考虑到可否在不合并的情况下，直接通过二分查找定位到合并后数组的中位数。

虽然算法不会合并nums1和nums2，但是依然会假设合并后经过排序的数组为nums3。

以0作为数组下标的初始下标，那么数组nums1由nums1[0..<m]组成，数组nums2由nums2[0..<n]组成，其中m和n是nums1和nums2的长度。

• 当m + n为偶数时，中位数为nums3[(m+n)/2-1]和nums3[(m+n)/2]的平均值。
• 当m + n为奇数时，中位数为nums3[(m+n-1)/2]。

因为不能直接合并nums3（这样会导致算法的复杂度变成O(N)），为了在不合并的情况下获取中位数，需要直接从nums1[0..<i]和nums2[0..<j]查找。

可以计算出，对于nums3，中位数必然存在于nums3[0..<(m+n)/2+1]内，其中(m+n)/2取的是下整。

因为nums1和nums2各自都是已排序的数组，可以证明nums3[0..<(m+n)/2+1]是由nums1[0..<i]和nums2[0..<j]组成，其中i + j = (m+n)/2+1，m >= i >= 0，m >= i >= 0 。

那么，寻找中位数的问题就被转换为寻找i和j的问题，且如果i已知，则j也已知，因为可以通过(m+n)/2+1-i获得。

那么接下去的问题就是，如何知道nums1[0..<i]和nums2[0..<j]组成的数组正是nums3[0..<(m+n)/2+1]。

分三种情况分析：

1. 当m=0时，nums1为空数组，不难得出当i=0且j=(m+n)/2+1时，nums1[0..<i]和nums2[0..<j]就是我们要找到子数组。
2. 当n=0时，同样可以得出当i=(m+n)/2+1且j=0时，nums1[0..<i]和nums2[0..<j]就是我们要找到子数组。
3. 当m>0且n>0时，当且仅当nums1[i-1]<=nums2[j-1]<=nums[i]，或者nums1[j-1]<=nums2[i-1]<=nums[j]时，nums1[0..<i]和nums2[0..<j]才是我们要找的子数组。

为了找到nums1[0..<i]和nums2[0..<j]，我们可以先任意指定一个i，因此j也得到了。然后我们通过以上分析法分析子数组是否找到，如果没找到，我们通过二分法移动i或者j，直到找到为止。

通过二分法移动i或者j时，移动的策略是：

1. 当nums1[i]<nums2[j-1]时，说明nums1还可以向右移，则需要右移i，左移j。
2. 当nums2[j]<nums2[i-1]时，说明nums2还可以向右移，则需要右移j，左移i。
3. 移动的“步数”取决于nums1和nums2移动范围，我们每次选择移动范围更小的那个进行二分，不然移动就会溢出，当确定了移动步数后，另一个数组只要朝相反的方向移动同样步数即可，这样i+j的结果始终恒定。

并且需要考虑如下边界情况：

1. 当i需要左移但是i=0时，说明nums1的数据全部小于nums2，直接返回nums2的对应子数组。
2. 当j需要左移但是j=0时，说明nums2的数据全部小于nums1，直接返回nums1的对应子数组。

当nums1[0..<i]和nums2[0..<j]顺利被找到时，接下去是解决如何得出中位数。
可以分如下情况：

1. 当m+n结果为奇数时，中位数等于nums1[i-1]，nums2[j-1]两个数（如果其中任何数因为边界问题不存在则忽略）中最大的值。
2. 当m+n结果为偶数时，中位数等于nums1[i-1]，nums1[i-2]，nums2[j-1]，nums2[j-2]四个数（如果因为边界问题不存在则忽略）中最大的两个值，取两个值的平均值。

基础思路说完，接下去就是上代码部分。

为了让算法思路更清晰，对代码进行适度的封装。首先因为涉及到“定制的”二分算法，所以封装一个结构体用来描述搜索范围。

struct RANGE {
    var begin = 0
    var end = 0
    var center = 0
    mutating func set(_ b: Int, _ e: Int, _ c: Int) {
        begin = b
        end = e
        center = c
    }
}

定义一个cmp函数，用来描述需要左移、需要右移还是已经找到子数组的情况。

func cmp(_ a: Int, _ b: Int) -> Int {
    guard a != -1 && b != -1 else { return 0 }
    if nums1[a] >= nums2[b] {
        return b + 1 >= nums2.count || nums2[b+1] >= nums1[a] ? 0 : 1
    }
    return a + 1 >= nums1.count || nums1[a+1] >= nums2[b] ? 0 : -1
}

定义一个mov函数，处理子数组移动的情况，当一个子数组范围往一边移动时，必然会让另一个子数组往相反方向移动。

func mov(_ r1: inout RANGE, _ r2: inout RANGE, _ mv: Int) {
    r1.end = r1.center
    r1.center -= mv
    r2.begin = r2.center + 1
    r2.center += mv
}

接下去就是最重要的搜索部分。搜索search()函数将返回一个范围值，返回的结果就是nums1和nums2各自的子数组结果。和之前描述的略有差异的是，返回的结果下标为子数组最后一个数字的下标。如果其中一个子数组不存在，则下标返回-1。

func search() -> (Int, Int) {
   let n = (nums1.count + nums2.count) / 2 + 1
   if nums1.isEmpty {
       return (-1, n - 1)
   } else if nums2.isEmpty {
       return (n-1, -1)
   }
   var range1 = RANGE()
   var range2 = RANGE()
   if nums1.count < nums2.count {
       range1.set(0, nums1.count, nums1.count / 2)
       range2.set(0, nums2.count, n - 2 - range1.center)
   } else {
       range2.set(0, nums2.count, nums2.count / 2)
       range1.set(0, nums1.count, n - 2 - range2.center)
   }
   var t = cmp(range1.center, range2.center)
   while t != 0 {
       if t > 0 {
           var mv = 0
           if range1.center - range1.begin < range2.end - 1 - range2.center {
               mv = (range1.center - range1.begin + 1) / 2
           } else {
               mv = (range2.end - range2.center) / 2
           }
           guard mv != 0 else { return (-1, n-1) }
           mov(&range1, &range2, mv)
       } else {
           var mv = 0
           if range1.end - 1 - range1.center < range2.center - range2.begin {
               mv = (range1.end - range1.center) / 2
           } else {
               mv = (range2.center - range2.begin + 1) / 2
           }
           if mv == 0 {
               return (n-1, -1)
           }
           mov(&range2, &range1, mv)
       }
       t = cmp(range1.center, range2.center)
   }
   return (range1.center, range2.center)
}

最后，函数answer用来计算已知返回了两个子数组的情况下，计算中位数的结果。

func answer(_ s: (Int, Int)) -> Double {
      var m1 = Int(Int32.min)
      var m2 = Int(Int32.min)
      if (nums1.count + nums2.count) % 2 == 0 {
          if s.0 != -1 {
              m1 = nums1[s.0]
              if s.0 > 0 { m2 = nums1[s.0-1] }
              if m2 > m1 { swap(&m1, &m2) }
          }
          if s.1 != -1 {
              if nums2[s.1] > m2 { m2 = nums2[s.1] }
              if m2 > m1 { swap(&m1, &m2) }
              if s.1 > 0 && nums2[s.1-1] > m2 { m2 = nums2[s.1-1] }
              if m2 > m1 { swap(&m1, &m2) }
          }
      } else {
          if s.0 != -1 { m1 = nums1[s.0] }
          if s.1 != -1 { m1 = max(m1, nums2[s.1]) }
          m2 = m1
      }
      return (Double(m1) + Double(m2)) / 2.0
  }
  let s = search()
  return answer(s)
}

整体代码如下：

class Solution {
    func findMedianSortedArrays(_ nums1: [Int], _ nums2: [Int]) -> Double {
        struct RANGE {
            var begin = 0
            var end = 0
            var center = 0
            mutating func set(_ b: Int, _ e: Int, _ c: Int) {
                begin = b
                end = e
                center = c
            }
        }
        func cmp(_ a: Int, _ b: Int) -> Int {
            guard a != -1 && b != -1 else { return 0 }
            if nums1[a] >= nums2[b] {
                return b + 1 >= nums2.count || nums2[b+1] >= nums1[a] ? 0 : 1
            }
            return a + 1 >= nums1.count || nums1[a+1] >= nums2[b] ? 0 : -1
        }
        func mov(_ r1: inout RANGE, _ r2: inout RANGE, _ mv: Int) {
            r1.end = r1.center
            r1.center -= mv
            r2.begin = r2.center + 1
            r2.center += mv
        }
        func search() -> (Int, Int) {
            let n = (nums1.count + nums2.count) / 2 + 1
            if nums1.isEmpty {
                return (-1, n - 1)
            } else if nums2.isEmpty {
                return (n-1, -1)
            }
            var range1 = RANGE()
            var range2 = RANGE()
            if nums1.count < nums2.count {
                range1.set(0, nums1.count, nums1.count / 2)
                range2.set(0, nums2.count, n - 2 - range1.center)
            } else {
                range2.set(0, nums2.count, nums2.count / 2)
                range1.set(0, nums1.count, n - 2 - range2.center)
            }
            var t = cmp(range1.center, range2.center)
            while t != 0 {
                if t > 0 {
                    var mv = 0
                    if range1.center - range1.begin < range2.end - 1 - range2.center {
                        mv = (range1.center - range1.begin + 1) / 2
                    } else {
                        mv = (range2.end - range2.center) / 2
                    }
                    guard mv != 0 else { return (-1, n-1) }
                    mov(&range1, &range2, mv)
                } else {
                    var mv = 0
                    if range1.end - 1 - range1.center < range2.center - range2.begin {
                        mv = (range1.end - range1.center) / 2
                    } else {
                        mv = (range2.center - range2.begin + 1) / 2
                    }
                    if mv == 0 {
                        return (n-1, -1)
                    }
                    mov(&range2, &range1, mv)
                }
                t = cmp(range1.center, range2.center)
            }
            return (range1.center, range2.center)
        }
        func answer(_ s: (Int, Int)) -> Double {
            var m1 = Int(Int32.min)
            var m2 = Int(Int32.min)
            if (nums1.count + nums2.count) % 2 == 0 {
                if s.0 != -1 {
                    m1 = nums1[s.0]
                    if s.0 > 0 { m2 = nums1[s.0-1] }
                    if m2 > m1 { swap(&m1, &m2) }
                }
                if s.1 != -1 {
                    if nums2[s.1] > m2 { m2 = nums2[s.1] }
                    if m2 > m1 { swap(&m1, &m2) }
                    if s.1 > 0 && nums2[s.1-1] > m2 { m2 = nums2[s.1-1] }
                    if m2 > m1 { swap(&m1, &m2) }
                }
            } else {
                if s.0 != -1 { m1 = nums1[s.0] }
                if s.1 != -1 { m1 = max(m1, nums2[s.1]) }
                m2 = m1
            }
            return (Double(m1) + Double(m2)) / 2.0
        }
        let s = search()
        return answer(s)
    }
}

解法 3：双数组二分查找 方法 2

双数组二分查找的方法2和方法1策略类似，唯一不同的是该方法定义了的搜索函数findValueAtIndex是 搜索nums1和nums2组成的nums3中的第k个元素 。而该函数的实现策略选用了双数组二分。

借助该函数，我们就能轻松的获得中位数。

源代码如下。

// O(log(n))复杂度解法 2
class Solution {
    class Scope {
        var begin = 0
        var end = 0
        var center: Int { !valid ? -1 : (begin + end) / 2 }
        var valid: Bool { end > begin }
        var count: Int { end - begin }
        func moveRight() { begin = center }
        func moveLeft() { if valid { end = center } }
    }
    
    func findMedianSortedArrays(_ nums1: [Int], _ nums2: [Int]) -> Double {
        func findValueAtIndex(_ index: Int) -> Int {
            let scope1 = Scope()
            let scope2 = Scope()
            scope1.end = nums1.count
            scope2.end = nums2.count
            
            while scope1.valid || scope2.valid {
                let v1 = scope1.valid ? nums1[scope1.center] : Int.min
                let v2 = scope2.valid ? nums2[scope2.center] : Int.min
                let currentIndex = scope1.center + scope2.center + 1
                
                if v1 < v2 {
                    if currentIndex == index {
                        if scope1.valid && scope1.center + 1 < scope1.end && nums1[scope1.center+1] < nums2[scope2.center] {
                            scope2.moveLeft()
                        } else {
                            return nums2[scope2.center]
                        }
                    } else if currentIndex < index {
                        scope1.count <= 1 ? scope2.moveRight() : scope1.moveRight()
                    } else if currentIndex > index {
                        scope2.moveLeft()
                    }
                } else {
                    if currentIndex == index {
                        if scope2.valid && scope2.center + 1 < scope2.end && nums2[scope2.center+1] < nums1[scope1.center] {
                            scope1.moveLeft()
                        } else {
                            return nums1[scope1.center]
                        }
                    } else if currentIndex < index {
                        scope2.count <= 1 ? scope1.moveRight() : scope2.moveRight()
                    } else {
                        scope1.moveLeft()
                    }
                }
            }
            return 0
        }
        
        let c1 = nums1.count
        let c2 = nums2.count
        if (c1 + c2) % 2 == 0 {
            let i1 = (c1 + c2) / 2
            let i2 = i1 - 1
            let v1 = findValueAtIndex(i1)
            let v2 = findValueAtIndex(i2)
            return (Double(v1) + Double(v2)) / 2.0
        } else {
            let i = (c1 + c2) / 2
            let value = findValueAtIndex(i)
            return Double(value)
        }
    }
}

总结

这是道个人认为非常经典考察对二分查找理解的题，该题目提供了两个已排序数组，如果要在O(1)空间和O(log(n))时间内解决该问题，就非要对两个数组进行二分查找而不进行重排不可。

力扣的题解里有非常精巧的代码解法，但是看了下基本思路无外乎我提供了第3和第4种解法，加上语言上的技巧可以让代码更精简。

关注我的公众号风海铜锣技术君点击加群按钮即可加入力扣解题交流群。

源代码可以到我的github上下载: https://github.com/FengHaiTongLuo/LeetCode4Swift/blob/main/4.%20Median%20of%20Two%20Sorted%20Arrays.swift