如何用Numpy高效构建由切片组成的不规则数组？

假设有一个大型numpy数组（约100万个浮点值）：

import numpy as np
data = np.arange(100).astype(np.float64)

同时定义两个整数数组，分别作为切片的起始和结束索引：

a = np.array([1,2,3,4,5,6,7])
b = np.array([10,20,30,40,50,60,70])

需要高效生成由data[a[i]:b[i]]组成的不规则数组，后续要对每个子数组求和，但np.take无法满足需求，求解决方案。

直接构建不规则数组（如numpy的object类型数组）会带来显著的内存开销和性能损耗，尤其针对百万级别的数据。更优的方案是利用前缀和数组直接计算切片和，完全无需生成子数组，实现全向量化操作。

方法1：前缀和快速计算切片和 #

步骤1：生成前缀和数组 #

前缀和数组prefix_sum中，prefix_sum[k]表示data[0:k]的累加和：

prefix_sum = np.concatenate([[0], np.cumsum(data)])

步骤2：向量化计算所有切片和 #

切片data[a[i]:b[i]]的和等价于prefix_sum[b[i]] - prefix_sum[a[i]]，直接批量计算：

slice_sums = prefix_sum[b] - prefix_sum[a]

验证示例：

• 第一个切片data[1:10]的和为1+2+...+9=45，计算得prefix_sum[10] - prefix_sum[1] = 45 - 0 = 45，结果正确。
• 第二个切片data[2:20]的和为2+3+...+19=189，计算得prefix_sum[20] - prefix_sum[2] = 190 - 1 = 189，结果正确。

该方法时间复杂度为O(n)，完全基于numpy向量化操作，无Python循环开销，是百万级数据场景下的最优选择。

方法2：生成子数组（仅适用于小数据量） #

如果确实需要生成子数组（不推荐用于百万级数据），可以使用列表推导：

subarrays = [data[start:end] for start, end in zip(a, b)]
# 对应求和操作
sub_sums = [np.sum(arr) for arr in subarrays]

但此方法会生成多个独立的子数组，内存占用高，且Python循环会降低处理速度，仅适合小规模数据场景。

内容的提问来源于stack exchange，提问作者A. G