使用索引列表对二维 np 数组的行求和答案

【问题标题】：Sum Rows of 2D np array with List of Indices使用索引列表对二维 np 数组的行求和
【发布时间】：2020-05-04 03:08:29
【问题描述】：

我有一个 2d numpy 数组和一个数字列表。如果列表是[1, 3, 1, 8]，其中列表总和为行数，我想输出一个数组，其中第一行不变，接下来的三行求和，第五行不变，其余八行求和。

举个例子：

A = [[0,0], [1,2], [3,4]] 和 l = [1, 2] 将输出 [[0,0], [4,6]

我查看了 np.sum 和其他功能，但找不到这个功能。谢谢。

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

您可以只遍历 l 的索引并根据位置获取该行或对一系列行求和。

import numpy as np

A = [[0,0], [1,2], [3,4]]
l = [1, 2]

ans = []
for i in range(len(l)):
  if i%2 == 0:
    ans.append(A[ l[i] ])
  else:
    ans.append( np.sum( A[ l[i-1]:l[i-1] + l[i] ], axis=0 ) )

ans = np.array(ans)
print(ans)

[[1 2]
 [4 6]]

注：

如果列表是 [1, 3, 1, 8]，其中列表总和为行数，我想输出一个第一行不变的数组，接下来的三个行总和，第五行不变，其余八行总和。

我想你的意思是[1, 3, 5, 8]

【讨论】：

【解决方案2】：

如果l中的元素数量比较多，使用pandas中的groupby可能会获得更好的性能，例如

import pandas as pd

labels = np.repeat(np.arange(1, len(l) + 1), l)
# [1, 2, 2]
df = pd.DataFrame(A)
df['label'] = labels
result = df.groupby('label').sum().values

【讨论】：

【解决方案3】：

当我意识到我可以在不影响所需输出的情况下对列表进行排序时，我最终提出了自己的解决方案。我使用 np.unique 来确定排序列表中每个元素的第一个索引，然后对这些索引之间的行求和。见下文。

elements, indices = np.unique(data, return_counts=True)
row_summing = np.append([0], np.cumsum(indices))[:-1] #[0, index1, index2,...]

output = np.add.reduceat(matrix, row_summing, axis=0)

【讨论】：

我认为这应该是公认的答案； reduceat 比 pandas 优雅得多。作为一个挑剔的人，您也许可以使用np.repeat 修改row_summing 定义以提高效率。由于散列与排序，pd.unique 比 np.unique 具有更好的渐近行为，这是众所周知的秘密，但这在这里可能不会发挥重要作用。