Numpy：条件总和答案

【问题标题】：Numpy: conditional sumNumpy：条件总和
【发布时间】：2018-05-01 18:39:23
【问题描述】：

我有以下 numpy 数组：

import numpy as np
arr = np.array([[1,2,3,4,2000],
                [5,6,7,8,2000],
                [9,0,1,2,2001],
                [3,4,5,6,2001],
                [7,8,9,0,2002],
                [1,2,3,4,2002],
                [5,6,7,8,2003],
                [9,0,1,2,2003]
              ])

我了解np.sum(arr, axis=0)提供结果：

array([   40,    28,    36,    34, 16012])

我想做的（没有 for 循环）是根据最后一列的值对列求和，以便提供的结果是：

array([[   6,    8,   10,   12, 4000],
       [  12,    4,    6,    8, 4002],
       [   8,   10,   12,    4, 4004],
       [  14,    6,    8,   10, 4006]])

我意识到如果没有循环可能会有些牵强，但希望能做到最好……

如果必须使用 for 循环，那将如何工作？

我试过np.sum(arr[:, 4]==2000, axis=0)（我会用for循环中的变量替换2000），但是它给出了2

的结果

【问题讨论】：

我认为您正在寻找pandas groupby。
右列中的值是否总是恰好重复两次，或者这只是您的示例中的巧合？
@cᴏʟᴅsᴘᴇᴇᴅ 你能重新打开吗？我正在研究一个纯 numpy 解决方案。
@MadPhysicist 好的，没问题，我也有兴趣看看。
巧合（我基本上有很多数据要按年汇总）。 df.groupby(4, axis=0).sum() 确实给了我我需要的东西。我不会回答，因为我想知道是否可以用 numpy 完成同样的事情，但是谢谢！

标签： python arrays numpy sum

【解决方案1】：

您可以使用 np.diff 和 np.add.reduceat 的巧妙应用在纯 numpy 中执行此操作。 np.diff 将为您提供最右侧列更改的索引：

d = np.diff(arr[:, -1])

np.where 会将您的布尔索引 d 转换为 np.add.reduceat 期望的整数索引：

d = np.where(d)[0]

reduceat 还期望看到零索引，并且所有内容都需要移动一：

indices = np.r_[0, e + 1]

在这里使用np.r_ 比np.concatenate 更方便一点，因为它允许使用标量。然后总和变为：

result = np.add.reduceat(arr, indices, axis=0)

这当然可以组合成一条线：

>>> result = np.add.reduceat(arr, np.r_[0, np.where(np.diff(arr[:, -1]))[0] + 1], axis=0)
>>> result
array([[   6,    8,   10,   12, 4000],
       [  12,    4,    6,    8, 4002],
       [   8,   10,   12,    4, 4004],
       [  14,    6,    8,   10, 4006]])

【讨论】：

不错的答案；尽管单行字很难阅读，但解释得很好:)
谢谢。我认为@Divakar 的回答是对同一想法的更强有力的诠释。
选择这个作为答案，因为它只使用numpy 回答问题，但说实话，我更喜欢@MadPhysicist 的pandas.groupby，它实际上对我的最终解决方案更有效，因为我也需要按月和年分组。谢谢大家。

【解决方案2】：

我发布了一个使用pandas 的简单解决方案和一个使用itertools 的解决方案

import pandas as pd
df = pd.DataFrame(arr)
x = df.groupby(4).sum().reset_index()[range(5)] #range(5) adjusts ordering 
x[4] *= 2
np.array(x)

array([[   6,    8,   10,   12, 4000],
       [  12,    4,    6,    8, 4002],
       [   8,   10,   12,    4, 4004],
       [  14,    6,    8,   10, 4006]])

你也可以使用itertools

np.array([sum(x[1]) for x in itertools.groupby(arr, key = lambda k: k[-1])])

array([[   6,    8,   10,   12, 4000],
       [  12,    4,    6,    8, 4002],
       [   8,   10,   12,    4, 4004],
       [  14,    6,    8,   10, 4006]])

【讨论】：

【解决方案3】：

方法 #1：基于 NumPy 的 sum-reduction

这是基于np.add.reduceat的一个-

def groupbycol(a, assume_sorted_col=False, colID=-1):
    if assume_sorted_col==0:
        # If a is not already sorted by that col, use argsort indices for
        # that colID and re-arrange rows accordingly
        sidx = a[:,colID].argsort()
        a_s = a[sidx] # sorted by colID col of input array
    else:
        a_s = a

    # Get group shifting indices
    cut_idx = np.flatnonzero(np.r_[True, a_s[1:,colID] != a_s[:-1,colID]])

    # Use those indices to setup sum reduction at intervals along first axis
    return np.add.reduceat(a_s, cut_idx, axis=0)

示例运行 -

In [64]: arr
Out[64]: 
array([[   1,    2,    3,    4, 2000],
       [   5,    6,    7,    8, 2000],
       [   9,    0,    1,    2, 2001],
       [   3,    4,    5,    6, 2001],
       [   7,    8,    9,    0, 2002],
       [   1,    2,    3,    4, 2002],
       [   5,    6,    7,    8, 2003],
       [   9,    0,    1,    2, 2003]])

In [65]: # Shuffle rows off input array to create a generic last col (not sorted)
    ...: np.random.seed(0)
    ...: np.random.shuffle(arr)

In [66]: arr
Out[66]: 
array([[   5,    6,    7,    8, 2003],
       [   9,    0,    1,    2, 2001],
       [   5,    6,    7,    8, 2000],
       [   9,    0,    1,    2, 2003],
       [   3,    4,    5,    6, 2001],
       [   1,    2,    3,    4, 2000],
       [   1,    2,    3,    4, 2002],
       [   7,    8,    9,    0, 2002]])

In [67]: groupbycol(arr, assume_sorted_col=False, colID=-1)
Out[67]: 
array([[   6,    8,   10,   12, 4000],
       [  12,    4,    6,    8, 4002],
       [   8,   10,   12,    4, 4004],
       [  14,    6,    8,   10, 4006]])

方法 #2：利用矩阵乘法

我们基本上可以用广播掩码创建 + 矩阵乘法替换 np.add.reduceat，因此利用快速 BLAS，它也适用于通用未排序列 -

import pandas as pd

def groupbycol_matmul(a, colID=-1):
    mask = pd.Series(a[:,colID]).unique()[:,None] == arr[:,colID]
    return mask.dot(arr)

【讨论】：

先做 argsort 很好。
希望我能给乘法另一个 +1。

【解决方案4】：

您可能想看看numpy_indexed。有了它，你可以做到：

import numpy as np
import numpy_indexed as npi

arr = np.array([[1,2,3,4,2000],
                [5,6,7,8,2000],
                [9,0,1,2,2001],
                [3,4,5,6,2001],
                [7,8,9,0,2002],
                [1,2,3,4,2002],
                [5,6,7,8,2003],
                [9,0,1,2,2003]
              ])


result = npi.GroupBy(arr[:, 4]).sum(arr)[1]

>>>[[   6    8   10   12 4000]
    [  12    4    6    8 4002]
    [   8   10   12    4 4004]
    [  14    6    8   10 4006]]

【讨论】：