考虑一个一维 NumPy 输入数组和一个排序索引数组。目标是获得输入数组 a
的总和,但按索引数组中定义的索引进行分割。
下面有两种方法,但它们都需要缓慢的 Python for 循环。是否有不需要 Python for 循环的纯 NumPy 版本?
示例:
a = np.arange(20) # Input array
idxs = np.array([7, 15, 16]) # Index array
# Goal: Split a at index 7, 15 and 16 and
# compute sum for each partition
# Solution 1:
idxs_ext = np.concatenate(([0], idxs, [a.size]))
results = np.empty(idxs.size + 1)
for i in range(results.size):
results[i] = a[idxs_ext[i]:idxs_ext[i+1]].sum()
# Solution 2:
result = np.array(
[a_.sum() for a_ in np.split(a, idxs)]
)
# Result: array([21., 84., 15., 70.])
最佳答案
首先,您可以根据 idxs
数组通过 np.split
拆分 a
数组,然后将函数应用到该数组上,如下所示:
np.stack(np.vectorize(np.sum)(np.array(np.split(a, idxs), dtype=object)))
另一个答案是使用 np.add.reduceat
正如 @hpaulj 在评论中提到的那样,速度更快:
np.add.reduceat(a, np.insert(idxs, 0, 0), axis=0)
更新:
对于具有 7 个切片的数据范围 1000,使用 np.concatenate
代替 insert
将运行时间减少了 5 倍; 这是我研究过的最快的方法:
np.add.reduceat(a, np.concatenate(([0], idxs)), axis=0)
关于python - NumPy:按索引分割的一维数组求和,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/70322757/