numpy数组：按一列分组，求和另一列

Question

我有一个看起来像这样的数组：

 array([[ 0,  1,  2],
        [ 1,  1,  6],
        [ 2,  2, 10],
        [ 3,  2, 14]])

我想对第二列中具有相同值的第三列的值求和，所以结果是：

 array([[ 0,  1,  8],
        [ 1,  2, 24]])

我开始对此进行编码，但是我坚持下面的总和：

import numpy as np
import sys

inFile = sys.argv[1]

with open(inFile, 'r') as t:
    f = np.genfromtxt(t, delimiter=None, names =["1","2","3"])

f.sort(order=["1","2"])
if value == previous.value:
   sum(f["3"])

Answer 1

如果您的数据按第二列排序，则可以使用以 为中心的内容np.add。reduceat纯 numpy 解决方案。应用np.nonzero( 或)的组合将为您提供第二列切换值的位置。您可以使用这些索引来进行求和。其他列非常公式化，因此您可以相当轻松地将它们连接回去：np.wherenp.diff

A = np.array([[ 0,  1,  2],
              [ 1,  1,  6],
              [ 2,  2, 10],
              [ 3,  2, 14]])
# Find the split indices
i = np.nonzero(np.diff(A[:, 1]))[0] + 1
i = np.insert(i, 0, 0)
# Compute the result columns
c0 = np.arange(i.size)
c1 = A[i, 1]
c2 = np.add.reduceat(A[:, 2], i)
# Concatenate the columns
result = np.c_[c0, c1, c2]

IDEOne链接

注意索引中的+1。这是因为考虑到工作原理，您总是想要切换之后的位置，而不是切换之前的位置reduceat。插入零作为第一个索引也可以使用np.r_、np.concatenate等来完成。

话虽这么说，我仍然认为您正在寻找@jpp's 答案中的 pandas 版本。

Answer 2

您可以使用pandas向量化算法：

import pandas as pd, numpy as np

A = np.array([[ 0,  1,  2],
              [ 1,  1,  6],
              [ 2,  2, 10],
              [ 3,  2, 14]])

df = pd.DataFrame(A)\
       .groupby(1, as_index=False)\
       .sum()\
       .reset_index()

res = df[['index', 1, 2]].values

结果

array([[ 0,  1,  8],
       [ 2,  2, 24]], dtype=int64)

Answer 3

使用字典存储值，然后转换回列表

x = [[ 0,  1,  2],
     [ 1,  1,  6],
     [ 2,  2, 10],
     [ 3,  2, 14]]

y = {}
for val in x:
    if val[1] in y:
        y[val[1]][2] += val[2]
    else:
        y.update({val[1]: val})
print([y[val] for val in y])