python实现连续数列相加,数列相加 - 问答 - Python中文网

循环numpy方法

这不是最漂亮的方法，而且可能有一种向量化方法，它更加优雅和高效，但是您可以使用列表理解来完成以下操作。基准见底部。在

给定一个二维数组x：>>> x

array([[ 1, 3, 7, 34],

[ 1, 4, 3, 6],

[ 1, 8, 2, 10],

[ 2, 9, 35, 11],

[ 2, 10, 12, 24]])

np.vstack([np.insert(x[x[:,0] == i][:,1:].sum(axis=0),0,i)

for i in np.unique(x[:,0])])

退货：

^{pr2}$

pandas建议

但是，对于这种表格数据，您可能需要考虑pandas。例如：import pandas as pd

df = pd.DataFrame(x)

>>> df

0 1 2 3

0 1 3 7 34

1 1 4 3 6

2 1 8 2 10

3 2 9 35 11

4 2 10 12 24

>>> df.groupby(0).sum()

1 2 3

0

1 15 12 50

2 19 47 35

您还可以很容易地将其恢复到numpy数组：>>> df.groupby(0).sum().reset_index().values

array([[ 1, 15, 12, 50],

[ 2, 19, 47, 35]])

基准

令我惊讶的是，这两种方法在速度上是可比的，我的loopy numpy方法险胜pandas方法：在形状(100000, 100)的大数组中，loopynumpy方法平均耗时约0.1秒，pandas方法平均耗时约0.13秒import timeit

x = np.random.randint(0,10,(100000, 100))

def np_method(arr = x):

return np.vstack([np.insert(arr[arr[:,0] == i][:,1:].sum(axis=0),0,i)

for i in np.unique(arr[:,0])])

def pd_method(df = pd.DataFrame(x)):

return df.groupby(0).sum().reset_index().values

>>> timeit.timeit(pd_method, number = 100) / 100

0.12673938989639283

>>> timeit.timeit(np_method, number = 100) / 100

0.09724574089050293