基于分治法与递归的两有序数组中位数求解及边界问题修复

需求说明 #

我需要用**分治法（divide and conquer）结合递归（recursion）**直接求解两个不同大小的有序数组的中位数，不允许先合并数组再计算。

示例：

• 给定ar1 = [1, 2, 3, 4, 12]和ar2 = [0, 7, 8, 11]，中位数为4
• 给定ar1 = [1, 2, 3, 4, 12]和ar2 = [0]，中位数为2

我已知多种解法，但只需要分治递归实现。当前代码能处理部分场景，但当其中一个数组为单元素时失效。我发现问题出在代码前两个if判断，但移除它们会导致递归无返回值。

原问题代码 #

import math

def medianOfTwo(L1, L2, l1, h1, l2, h2):

  if (h1 == l1): return L1[l1]
  if (h2 == l2): return L2[l2]

  medL1 = math.floor((l1+h1)/2)
  medL2 = math.floor((l2+h2)/2)


  if(L1[medL1] == L2[medL2]):
    return L1[medL1]
  

  if(L1[medL1] < L2[medL2]):
    return medianOfTwo(L1, L2, medL1, h1, l2, medL2)
  else:
    return medianOfTwo(L1, L2, l1, medL1, medL2, h2)


#ar1 = [1, 2, 3, 4, 12] 
#ar2 = [0, 7, 8, 11] 

ar1 = [1, 2, 3, 4] 
#ar2 = [7, 8, 11, 12]
ar2 = [0]

print(medianOfTwo(ar1, ar2, 0, len(ar1)-1, 0, len(ar2)-1))

问题分析 #

原代码的核心错误在于：当其中一个数组只剩单元素时，直接返回该元素作为中位数，忽略了中位数是两个数组合并后的中间值这一本质。比如测试用例ar1=[1,2,3,4]、ar2=[0]，合并后数组是[0,1,2,3,4]，中位数应为2，但原代码直接返回ar2[0]=0，完全不符合预期。

同时，原分治逻辑的区间裁剪规则不合理，可能导致递归无法收敛。

修复后的代码 #

import math

def medianOfTwo(L1, L2, l1, h1, l2, h2):
    # 计算当前子数组长度
    len1 = h1 - l1 + 1
    len2 = h2 - l2 + 1

    # 边界情况1：L1只剩一个元素
    if len1 == 1:
        # 找到该元素在L2中的插入位置
        insert_pos = l2
        while insert_pos <= h2 and L2[insert_pos] < L1[l1]:
            insert_pos += 1
        total_len = len1 + len2
        mid_idx = (total_len - 1) // 2
        # 根据插入位置计算中位数
        if insert_pos == l2:
            return L2[l2 + mid_idx] if mid_idx < len2 else L1[l1]
        elif insert_pos == h2 + 1:
            return L2[l2 + mid_idx - 1] if mid_idx >= 1 else L1[l1]
        else:
            if mid_idx == insert_pos - l2:
                return L1[l1]
            elif mid_idx < insert_pos - l2:
                return L2[l2 + mid_idx]
            else:
                return L2[l2 + mid_idx - 1]
    
    # 边界情况2：L2只剩一个元素，逻辑和L1单元素对称
    if len2 == 1:
        insert_pos = l1
        while insert_pos <= h1 and L1[insert_pos] < L2[l2]:
            insert_pos += 1
        total_len = len1 + len2
        mid_idx = (total_len - 1) // 2
        if insert_pos == l1:
            return L1[l1 + mid_idx] if mid_idx < len1 else L2[l2]
        elif insert_pos == h1 + 1:
            return L1[l1 + mid_idx - 1] if mid_idx >= 1 else L2[l2]
        else:
            if mid_idx == insert_pos - l1:
                return L2[l2]
            elif mid_idx < insert_pos - l1:
                return L1[l1 + mid_idx]
            else:
                return L1[l1 + mid_idx - 1]

    # 分治核心逻辑：计算当前子数组的中位数索引
    medL1 = (l1 + h1) // 2
    medL2 = (l2 + h2) // 2

    if L1[medL1] == L2[medL2]:
        return L1[medL1]
    
    # 每次递归裁剪相同数量的元素，保证递归收敛
    remove_count = min(len1 // 2, len2 // 2)
    if L1[medL1] < L2[medL2]:
        # 保留L1中位数右侧，L2中位数左侧（裁剪相同数量元素）
        return medianOfTwo(L1, L2, medL1, h1, l2, h2 - remove_count)
    else:
        # 保留L1中位数左侧，L2中位数右侧（裁剪相同数量元素）
        return medianOfTwo(L1, L2, l1, h1 - remove_count, medL2, h2)


# 测试示例1
ar1 = [1, 2, 3, 4, 12] 
ar2 = [0, 7, 8, 11] 
print(medianOfTwo(ar1, ar2, 0, len(ar1)-1, 0, len(ar2)-1))  # 输出4

# 测试示例2
ar1 = [1, 2, 3, 4, 12] 
ar2 = [0]
print(medianOfTwo(ar1, ar2, 0, len(ar1)-1, 0, len(ar2)-1))  # 输出2

# 测试示例3
ar1 = [1, 2, 3, 4] 
ar2 = [0]
print(medianOfTwo(ar1, ar2, 0, len(ar1)-1, 0, len(ar2)-1))  # 输出2

修复说明 #

1. 修正单元素边界逻辑：当某数组只剩单元素时，不再直接返回该元素，而是计算其在另一个数组中的插入位置，结合总长度推导合并后的中位数。
2. 优化分治裁剪规则：每次递归裁剪两个数组中长度较小的那一半的元素数量，确保递归能够收敛，避免无限递归。
3. 覆盖所有场景：包含不同长度数组、单元素数组等边界情况的正确处理。

内容的提问来源于stack exchange，提问作者Catarina Nogueira