[英]NumPy: Sum over 1-D array split by index
考慮一維NumPy輸入數組和排序索引數組。 目標是獲得輸入數組a
的總和,但由索引數組中定義的索引拆分。
以下是兩種方法,但它們都需要慢速 Python 循環。 是否有不需要 Python 循環的純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.])
首先,您可以通過np.split
根據您的idxs
數組拆分a
數組,然后將 function 應用為:
np.stack(np.vectorize(np.sum)(np.array(np.split(a, idxs), dtype=object)))
另一個答案是使用np.add.reduceat
在評論中提到的 np.add.reduceat 並且更快:
np.add.reduceat(a, np.insert(idxs, 0, 0), axis=0)
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.