求满足双条件的数组平衡拆分的O(n)时间复杂度优化方案

问题描述 #

给定一个整数数组nums，编写函数判断能否将其拆分为两个数组A和B，满足以下两个条件：

1. A和B的元素总和相等；
2. A中的所有元素严格小于B中的所有元素。

示例 #

• 输入nums = [1,5,7,1] → 返回true，拆分后A = [1,1,5]，B = [7]，满足两个条件；
• 输入nums = [12,7,6,7,6] → 返回false，即使拆分出和相等的A=[6,6,7]、B=[7,12]，但A中存在元素7与B中的7相等，不满足严格小于的要求。

我的现有解法（O(nlogn)时间复杂度） #

我通过排序数组实现了解法，代码如下：

from typing import List


def solution(nums: List[int]) -> bool:
    total_sum = sum(nums)

    # 总和为0或奇数时，无法拆分成和相等的两个数组
    if total_sum % 2 or total_sum == 0:
        return False

    nums.sort()

    curr_sum, i = total_sum, 0

    while curr_sum > total_sum // 2:
        curr_sum -= nums[i]
        i += 1

    # 当当前和等于目标和，且拆分点的元素不等于前一个元素时，满足条件
    if curr_sum == total_sum // 2 and nums[i] != nums[i - 1]:
        return True

    return False

是否存在O(n)时间复杂度的解法？ #

答案是分情况讨论：

情况1：数组元素取值范围有限 #

如果数组中的元素取值在固定的小范围内（比如0~1000），可以利用计数排序的思想实现严格O(n)的解法，步骤如下：

1. 计算数组总和，若为奇数或0直接返回false；
2. 统计每个数字的出现次数（用数组或哈希表）；
3. 从最小的数字开始累加元素和，直到累加和达到总和的一半；
4. 此时检查是否存在比当前累加的最大数字更大的元素——如果存在，说明可以将所有小于等于该数字的元素归入A，更大的归入B，满足严格小于的条件，返回true；否则返回false。

对应的Python实现代码：

from typing import List
import collections

def solution_O_n(nums: List[int]) -> bool:
    total_sum = sum(nums)
    if total_sum % 2 != 0 or total_sum == 0:
        return False
    target = total_sum // 2

    count = collections.Counter(nums)
    min_num = min(nums)
    max_num = max(nums)

    current_sum = 0

    for num in range(min_num, max_num + 1):
        if num not in count:
            continue
        # 尽可能多地累加当前数字，直到接近目标和
        while count[num] > 0 and current_sum + num <= target:
            current_sum += num
            count[num] -= 1
        if current_sum == target:
            # 检查是否存在比当前数字更大的元素
            return num < max_num
        elif current_sum > target:
            return False

    return False

情况2：数组元素取值范围无界 #

如果数组元素可以是任意大的整数，那么基于比较的排序无法突破O(nlogn)的时间下界，此时不存在严格意义上的O(n)解法。不过可以通过统计唯一元素的频率，再遍历唯一元素（若唯一元素数量远小于n）来优化实际运行效率，但本质上时间复杂度仍取决于唯一元素的排序时间。

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