python - numpy 对几列求和

Question

我想读取具有以下结构的文件:

1         3      7        34
1         4      3         6
1         8      2        10
2         9     35        11
2        10     12        24
...

列数从一个文件更改为另一个文件，我想将 1、2 等的所有列值相加。

所以结果应该是这样的

1       15     12      50
2       19     37      45
...

我已经开始编写类似的代码:

import numpy as np
import sys

inFile = sys.argv[1]

A = np.genfromtxt(inFile, delimiter=None)

a = np.size(A,1) #get the number of columns

i = np.nonzero(np.diff(A[:, 0]))[0] + 1
i = np.insert(i, 0, 0)

c1 = A[i, 0]
for l in range(3,a+1):
    c = np.add.reduceat(A[:, l], i)
    result = np.c_[c1, c]

如何使用 numpy 做到这一点？

Answer 1

Loopy 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])])

返回:

array([[ 1, 15, 12, 50],
       [ 2, 19, 47, 35]])

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]])

基准

令我惊讶的是，这两种方法就速度而言是相当的，我的循环 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