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

- Anom

---

澄清一下，看起来你的第一列是行号索引，与你的数据无关。然后，您希望第二列是该列中唯一的元素集，第三列是这些集合元素的现有第三列的总和。您的数据是否已按第二列排序，就像在您的示例中一样？ - Scott Mermelstein

是的，这是排序后的数据。我只是添加了第一列作为指示，表明我有“无用”信息的列。 - Anom

如果你省略那一列，实际上会更容易。 - JahKnows

你考虑过使用pandas吗？它可以为你生成索引列并进行分组。 - Mad Physicist

@Anom，如果以下任何一种解决方案有所帮助，请考虑接受它（左侧的绿色勾号），以便其他用户知道。 - jpp

Answer 1

如果您的数据是按第二列排序的，您可以使用围绕np.add和reduceat的东西来实现一个纯numpy解决方案。将np.diff应用于np.nonzero（或np.where）可以得到第二列切换值的位置。您可以使用这些索引进行总和缩减。其他列相当公式化，因此您可以相对容易地将它们连接回去：

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的工作方式。将零插入到第一个索引的操作也可以使用{{link2：np.r_}}, np.concatenate等方法实现。

话虽如此，我仍然认为你正在寻找@jpp's答案中的pandas版本。

- Mad Physicist

---

你是否在某个地方传递了一个一维数组而没有意识到？ - Mad Physicist

关于这个解决方案与 pandas.groupby 的比较，这个纯 numpy 解决方案要快得多。 - Ketil Tveiten

@Ketil。Numpy通常倾向于更快一些。但对于更复杂的问题来说，它也不太易读和使用。 - Mad Physicist

@MadPhysicist 当然可以；对于那些想要加速的人来说，了解如何避免使用 pandas 可能会很有用。虽然可能不太美观，但它确实很快！ - Ketil Tveiten

@MadPhysicist 这非常有帮助，因为我需要在工作中使用纯NumPy解决类似的问题。后来在晚餐时，我意识到我的问题可以用一个基于bin计数的一行代码解决。再次审查OP的问题后，我相信这里可能适用hist方法。很快会发布一个答案。 - Mercury

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)

- jpp

---

我的代码则完全使用了numpy。主要是为了说服楼主选择pandas :) - Mad Physicist

此外，我相信 OP 不是在寻找 0: 'first'。第一列是 0, 1，而不是 OP 期望结果中的 0, 2。 - Mad Physicist

Answer 3

使用 np.histogram 可以实现非常整洁、纯正的 numpy 解决方案：

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

c1 = np.unique(A[:, 1])
c0 = np.arange(c1.shape[0])
c2 = np.histogram(A[:, 1], weights=A[:, 2], bins=c1.shape[0])[0]

result = np.c_[c0, c1, c2]

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

当在np.histogram中提供一个与输入数组相同形状的weights数组时，输入数组a中的任意元素a[i]将会为其所属的箱子的计数贡献weights[i]。

例如，我们正在对第二列进行计数，而不是计算2出现了2次，我们将得到2出现了10次 + 2出现了14次 = 在2的箱子中的数量为28。

- Mercury

Answer 4

这是我的解决方案，只使用了numpy数组...

import numpy as np
arr = np.array([[ 0,  1,  2], [ 1,  1,  6], [ 2,  2, 10], [ 3,  2, 14]])

lst = []
compt = 0
for index in range(1, max(arr[:, 1]) + 1):
    lst.append([compt, index, np.sum(arr[arr[:, 1] == index][:, 2])])
lst = np.array(lst)
print lst
# lst, outputs...
# [[ 0  1  8]
# [ 0  2 24]]

这里的难点是np.sum(arr[arr[:, 1] == index][:, 2])，所以让我们将其分解为多个部分。

• arr[arr[:, 1] == index]的意思是...

你有一个数组arr，我们要求numpy匹配循环变量值的行。在这里，它从1设置到第二列（即索引为1）的元素的最大值。仅打印循环中的此表达式会得到...

# First iteration
[[0 1 2]
 [1 1 6]]
# Second iteration
[[ 2  2 10]
 [ 3  2 14]]

• 在我们的表达式中添加[:, 2]，意味着我们想要上述列表的第3列的值（即索引为2）。如果我打印arr[arr[:, 1] == index][:, 2]，它会给我...在第一次迭代时是[2, 6]，在第二次迭代时是[10, 14]。

• 我只需要使用np.sum()对这些值进行求和，并相应地格式化我的输出列表。 :)

- IMCoins

---

你正在寻找ufunc.reduceat。这可以在不使用循环的情况下完成。请查看我的答案。 - Mad Physicist

@MadPhysicist 嗯，这就是我来到这个网站的原因：向他人学习。我正在看一下，谢谢。 :) - IMCoins

我记得有人第一次向我展示reduceat。它很神秘，但非常方便。通常，如果您正在使用numpy数组的循环，您可能需要重新考虑您的方法。 - Mad Physicist

@MadPhysicist 我已经听说过这个观点，并且完全同意。但是，有时候在这种情况下，我没有整天的时间去寻找一个能让我的工作更轻松的特殊函数。:'( -- 话虽如此，一开始我正在寻找 np.where()，但是无法弄清楚如何使用它。现在我也知道我应该将其与 np.diff() 结合使用。:p - IMCoins

Answer 5

使用字典来存储值，然后将其转换回列表。

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

- JahKnows

Answer 6

要获得精确的输出，请使用pandas：

import pandas as pd
import numpy as np

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

df = pd.DataFrame(a)
df.groupby(1).sum().reset_index().reset_index().as_matrix()
#[[ 0 1  8]
# [ 1 2 24]]

- zipa

---

reset_index().reset_index()? - Mad Physicist

此外，您可能希望执行 reset_index(inplace=True)。 - Mad Physicist

@MadPhysicist 第一个重置了 groupby，第二个添加了包含索引值的新列，并匹配所需的输出。 - zipa

结果并不是你所声称的那样...请看一下你的第二列和jpp的答案。 - Mad Physicist

Answer 7

你也可以使用defaultdict并对值进行求和：

from collections import defaultdict

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

res = defaultdict(int)
for val in x:
    res[val[1]]+= val[2]
print ([[i, val,res[val]] for i, val in enumerate(res)])

- Hirabayashi Taro

---

我认为这不能保证原始数组的顺序（因为字典是无序的）。 - ChatterOne

我也是这么想的，而且我很惊讶在Python 3中使用正整数键时，我总是得到一个排序后的结果。 - Hirabayashi Taro

最近的Python版本中，字典现在是有序的。 - hpaulj

list(res.items()) 可以替换最后一个语句。 - hpaulj