如何用更Pythonic的方法查找列表中总和等于总合的子序列或元素组合

接收一个列表后遍历，检查是否存在某段连续元素序列的总和等于整个列表的总和，若存在则返回对应的子列表，也可提供支持检查任意元素组合是否满足求和条件的方案。

• 输入1：example1 = [8, 9, 10, 10, 10, -20]
  输出：[8, 9, 10]
• 输入2：example2 = [-15, 6, 8, 2, 10, 10, -5]
  输出：[6, 8, 2]

方案1：连续子序列查找（前缀和优化） #

该方案是最优解，时间复杂度O(n)、空间复杂度O(n)，仅需一次遍历即可得到结果，符合Pythonic简洁高效的编码风格：

def find_continuous_subarray(nums: list[int]) -> list[int]:
    total = sum(nums)
    # 存储前缀和对应的最早下标，初始值处理前缀和直接等于总和的场景
    prefix_map = {0: -1}
    current_sum = 0
    for idx, num in enumerate(nums):
        current_sum += num
        # 存在符合要求的子序列，截取返回
        if current_sum - total in prefix_map:
            start_idx = prefix_map[current_sum - total] + 1
            return nums[start_idx: idx + 1]
        prefix_map[current_sum] = idx
    # 无符合条件的子序列返回空列表
    return []

# 测试代码
if __name__ == "__main__":
    example1 = [8, 9, 10, 10, 10, -20]
    print(find_continuous_subarray(example1))  # 输出 [8, 9, 10]
    example2 = [-15, 6, 8, 2, 10, 10, -5]
    print(find_continuous_subarray(example2))  # 输出 [6, 8, 2]

方案2：任意元素组合查找（动态规划法） #

如果不需要子序列连续，仅要求元素组合求和等于列表总和，可以用下面的动态规划方案快速找到第一个符合要求的组合：

def find_any_subset(nums: list[int]) -> list[int]:
    total = sum(nums)
    # dp存储和为key的对应元素组合，初始值处理单个元素等于总和的场景
    dp = {0: []}
    for num in nums:
        # 遍历当前已有的和值，避免字典修改导致的遍历异常
        for exist_sum in list(dp.keys()):
            new_sum = exist_sum + num
            if new_sum == total:
                return dp[exist_sum] + [num]
            if new_sum not in dp and new_sum <= total:
                dp[new_sum] = dp[exist_sum] + [num]
    # 无符合条件的组合返回空列表
    return []

如果需要获取所有符合条件的组合，修改代码收集所有匹配的路径即可。

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