如何编写通用nSum函数找出数组中和为0的元素？

我需要编写一个程序，输入不含0的正负整数数组，返回其中和为0的元素组合；若无符合条件的元素则返回“No Elements found”。目前已分别实现了twoSum、threeSum、fourSum等函数，现在需要编写一个通用的nSum函数（n为待选元素的个数）。例如在数组[-4, -3, -2, -1, 10]中，应返回这5个元素，因其和为0。

以下是我已实现的代码：

def twoSum(arr, n, x, lp, rp):
    while lp < rp:
        if arr[lp] + arr[rp] == x:
            return True, [arr[lp], arr[rp]]
        elif arr[lp] + arr[rp] < x:
            lp += 1
        else:
            rp -= 1
    return False, []

def threeSum(arr, n, x, firstIndex):
    for i in range(firstIndex, n-2):
        check, res = twoSum(arr, n, x - arr[i], i+1, n-1)
        if check:
            return True, [arr[i]] + res
    return False, []
            
def finding_numbers(arr, n):
    if arr[0] < 0:
        check, res = twoSum(arr, n, 0, 0, n-1)
        if check:
            return res
        
        check, res = threeSum(arr, n, 0, 0)
        if check:
            return res
          
        for i in range(n-3):
            check, res = threeSum(arr, n, -arr[i], i+1)
            if check:
                return [arr[i]] + res
    return 'No Elements found'

# 调用示例（需先定义array和n）
# finding_numbers(array, n)

通用nSum函数实现方案 #

要实现通用的nSum，可通过递归抽象逻辑，核心思路如下：

• 当n == 2时，直接用双指针法求解（复用twoSum的逻辑）
• 当n > 2时，遍历数组中的每个元素，递归调用(n-1)Sum，将目标和调整为目标和 - 当前元素，同时起始索引设为当前索引+1，避免重复选取元素

完整的通用实现代码如下：

def n_sum(arr, target, n, start):
    length = len(arr)
    result = []
    # 递归终止条件：n为2时用双指针
    if n == 2:
        left, right = start, length - 1
        while left < right:
            current_sum = arr[left] + arr[right]
            if current_sum == target:
                result.append([arr[left], arr[right]])
                # 跳过重复元素（数组有重复时可选）
                while left < right and arr[left] == arr[left+1]:
                    left += 1
                while left < right and arr[right] == arr[right-1]:
                    right -= 1
                left += 1
                right -= 1
            elif current_sum < target:
                left += 1
            else:
                right -= 1
        return result
    # n>2时递归处理
    else:
        for i in range(start, length - n + 1):
            # 跳过重复元素（可选）
            if i > start and arr[i] == arr[i-1]:
                continue
            # 递归调用(n-1)Sum，目标和为target - arr[i]，起始索引为i+1
            sub_results = n_sum(arr, target - arr[i], n-1, i+1)
            for sub_res in sub_results:
                result.append([arr[i]] + sub_res)
        return result

def find_zero_sum_elements(arr):
    # 双指针法依赖有序数组，先排序
    sorted_arr = sorted(arr)
    length = len(sorted_arr)
    # 从n=2到n=数组长度依次尝试找和为0的组合
    for n in range(2, length + 1):
        res = n_sum(sorted_arr, 0, n, 0)
        if res:
            # 返回第一个找到的组合，若要返回所有组合可去掉[0]
            return res[0]
    return 'No Elements found'

# 调用示例
test_array = [-4, -3, -2, -1, 10]
print(find_zero_sum_elements(test_array))  # 输出: [-4, -3, -2, -1, 10]

说明 #

1. 通用nSum函数支持任意n>=2的情况，依赖数组有序，因此顶层函数先对输入数组排序
2. 代码中加入了跳过重复元素的逻辑（可选），若输入数组有重复元素，可避免返回重复组合
3. 顶层函数find_zero_sum_elements从n=2开始尝试，直到n等于数组长度，找到第一个符合条件的组合就返回

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