如何优化NumPy大规模数组的搜索与合并？能否用GPU加速？

问题背景 #

我有三个数组：存储ID的arr数组、指定区间的range数组，以及规模极大的query数组。query中的值部分存在于arr中，部分不存在；同时部分值落在range指定的区间内，部分不在。需要按规则获取数据：query中存在于arr的值从源A取数，不存在的从源B取数，最终结果要保持query的原始顺序。当前实现因为数组规模太大运行缓慢，想找性能优化方法，同时询问是否可用GPU加速。

原实现代码：

import numpy as np

arr = np.array([1,7,3,9,5,10,2,8,4,6])
range_arr = np.array([2,6])  # 重命名避免和Python关键字冲突
query = np.array([3,15,6,8,13,5,19])

contains     = []; contains_idx     = []
not_contains = []; not_contains_idx = []
in_range     = []; in_range_idx     = []
not_in_range = []; not_in_range_idx = []
contains_mask = np.full(len(query), False)
range_mask = np.full(len(query), False)

for idx,x in enumerate(query):
    if x in arr:
        contains.append(x)
        contains_idx.append(idx)
        contains_mask[idx] = True
    else:
        not_contains.append(x)
        not_contains_idx.append(idx)
    if x >= range_arr[0] and x <= range_arr[1]:
        in_range.append(x)
        in_range_idx.append(idx)
        range_mask[idx] = True
    else:
        not_in_range.append(x)
        not_in_range_idx.append(idx)

result = np.zeros((len(query), 5)) # 5维度n×5数组
# 简化处理，用随机数模拟源A、源B的数据
data_for_contains = np.random.random((len(contains), 5))
data_for_not_contains = np.random.random((len(not_contains), 5))
# 按query顺序填充结果
result[contains_idx] = data_for_contains
result[not_contains_idx] = data_for_not_contains
# 区间相关的操作逻辑类似，此处省略

性能优化方案 #

1. 用集合优化存在性判断 #

原代码中x in arr是O(n)时间复杂度，对大规模数组极慢。把arr转成Python集合后，x in set是O(1)，能大幅降低时间开销：

arr_set = set(arr)

2. 向量化操作替代Python循环 #

利用Numpy的向量化API直接生成掩码，完全避免循环，这是提升大规模数据处理速度的核心：

# 生成存在性掩码
contains_mask = np.isin(query, arr)
# 生成区间掩码
range_mask = (query >= range_arr[0]) & (query <= range_arr[1])

3. 省略不必要的中间列表 #

不需要单独存储contains、contains_idx等列表，直接用掩码索引操作结果数组，节省内存和数据复制开销：

result = np.zeros((len(query), 5))
# 填充源A数据
result[contains_mask] = np.random.random((contains_mask.sum(), 5))  # 模拟源A数据
# 填充源B数据
result[~contains_mask] = np.random.random(((~contains_mask).sum(), 5))  # 模拟源B数据

优化后的完整代码：

import numpy as np

arr = np.array([1,7,3,9,5,10,2,8,4,6])
range_arr = np.array([2,6])
query = np.array([3,15,6,8,13,5,19])

# 核心优化：集合+向量化操作
arr_set = set(arr)
contains_mask = np.isin(query, arr)
range_mask = (query >= range_arr[0]) & (query <= range_arr[1])

# 生成结果
result = np.zeros((len(query), 5))
# 模拟从源A、源B取数
result[contains_mask] = np.random.random((contains_mask.sum(), 5))
result[~contains_mask] = np.random.random(((~contains_mask).sum(), 5))

# 区间相关操作同理，直接用range_mask索引即可

GPU加速可行性 #

完全可以用GPU加速，适合处理超大规模数据（千万级以上），核心思路是用GPU张量库替代Numpy：

• 推荐工具：使用CuPy（API与Numpy完全兼容，无需大幅修改代码），或者TensorFlow/PyTorch的张量操作。
• 注意事项：GPU加速的优势体现在数据规模足够大时，小数据场景下，CPU到GPU的数据传输开销可能抵消并行计算的收益。
• CuPy示例：

import cupy as cp

# 转成CuPy张量
arr_cp = cp.array(arr)
range_cp = cp.array(range_arr)
query_cp = cp.array(query)

# 向量化操作（和Numpy完全一致）
contains_mask_cp = cp.isin(query_cp, arr_cp)
range_mask_cp = (query_cp >= range_cp[0]) & (query_cp <= range_cp[1])

# 生成结果
result_cp = cp.zeros((len(query_cp), 5))
result_cp[contains_mask_cp] = cp.random.random((contains_mask_cp.sum(), 5))
result_cp[~contains_mask_cp] = cp.random.random(((~contains_mask_cp).sum(), 5))

# 转回Numpy数组（如果需要）
result = cp.asnumpy(result_cp)

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