Python3下优化升序整数数组两数求和检测函数的极致性能方案

问题背景 #

现有如下函数代码，功能是接收一个由超大整数组成的数组，判断数组中是否存在两个不同元素之和等于指定Number并返回布尔值。由于该函数会被极多次调用，需要尽可能优化运行效率，替代嵌套循环实现：

flag = 0
for idx, i in enumerate(arr):
    for j in range(idx + 1, len(arr)):
        if i + arr[j] == Number:
            flag = 1
            break
    break
if flag == 0:
    return False
return True

补充信息（2023年8月13日） #

• 函数会被极多次调用
• Number和arr每次调用都会变化：Number取值范围极广，arr大小通常小于sqrt(Number)且按升序排列
• 希望在Kaggle Kernels中充分利用CPU/GPU资源提升性能

高效替代方案 #

1. 双指针法（排序数组最优解） #

由于数组已按升序排列，双指针法是当前场景下的最优选择，时间复杂度为O(n)，空间复杂度O(1)，完全适配高频调用需求：

def check_two_sum(arr, Number):
    left = 0
    right = len(arr) - 1
    while left < right:
        current_sum = arr[left] + arr[right]
        if current_sum == Number:
            return True
        elif current_sum < Number:
            left += 1
        else:
            right -= 1
    return False

核心优势：

• 无额外内存开销，内存占用极低
• 遍历次数最少，时间开销远低于嵌套循环的O(n²)
• 逻辑简单，CPU缓存友好，高频调用下性能表现稳定

2. 集合查找法（兼容通用场景） #

如果需要兼容未排序数组的场景，集合查找法时间复杂度O(n)，空间复杂度O(n)，但在排序数组场景下效率略逊于双指针法：

def check_two_sum(arr, Number):
    seen = set()
    for num in arr:
        complement = Number - num
        if complement in seen:
            return True
        seen.add(num)
    return False

3. 向量化/并行优化（适配Kaggle硬件资源） #

针对Kaggle的CPU/GPU资源，可以利用向量化操作或GPU加速进一步提升批量调用的吞吐量：

NumPy向量化（CPU加速）

import numpy as np

def check_two_sum_vectorized(arr, Number):
    arr_np = np.array(arr)
    # 生成所有两数之和的矩阵，排除自身相加的情况
    sum_matrix = arr_np[:, None] + arr_np[None, :]
    np.fill_diagonal(sum_matrix, -1)  # 用不可能等于Number的值覆盖对角线
    return np.any(sum_matrix == Number)

注意：此方法适合小尺寸数组，若数组较大，优先选择双指针法的向量化实现。

CuPy GPU加速

import cupy as cp

def check_two_sum_gpu(arr, Number):
    arr_cp = cp.array(arr)
    left = 0
    right = len(arr_cp) - 1
    while left < right:
        current_sum = arr_cp[left] + arr_cp[right]
        if current_sum == Number:
            return True
        elif current_sum < Number:
            left += 1
        else:
            right -= 1
    return False

说明：GPU加速适合处理极大数组或批量调用场景，需注意数据在CPU与GPU间的传输开销，避免抵消加速收益。

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