聚合数组的pythonic方法（numpy与否）

Question

我想做一个很好的函数来聚合数组中的数据（它是一个 numpy 记录数组，但它不会改变任何东西）

您有一个要在一个轴之间聚合的数据数组：例如一个dtype=[(name, (np.str_,8), (job, (np.str_,8), (income, np.uint32)]并且您希望获得每份工作的平均收入

我做了这个功能，在示例中它应该被称为aggregate(data,'job','income',mean)

---

def aggregate(data, key, value, func):

    data_per_key = {}

    for k,v in zip(data[key], data[value]):

        if k not in data_per_key.keys():

            data_per_key[k]=[]

        data_per_key[k].append(v)

    return [(k,func(data_per_key[k])) for k in data_per_key.keys()]

---

问题是我觉得它不是很好我想把它放在一行中：你有什么想法吗？

谢谢你的回答路易斯

PS：我想在通话中保留 func 以便您也可以要求中位数，最小值......

Answer 1

您的if k not in data_per_key.keys()可以被重写，就if k not in data_per_key ，但是您可以使用defaultdict做得更好。 这是一个使用defaultdict摆脱存在检查的版本：

import collections

def aggregate(data, key, value, func):
    data_per_key = collections.defaultdict(list)
    for k,v in zip(data[key], data[value]):
        data_per_key[k].append(v)

    return [(k,func(data_per_key[k])) for k in data_per_key.keys()]

Answer 2

也许您正在寻找的功能是matplotlib.mlab.rec_groupby ：

import matplotlib.mlab

data=np.array(
    [('Aaron','Digger',1),
     ('Bill','Planter',2),
     ('Carl','Waterer',3),
     ('Darlene','Planter',3),
     ('Earl','Digger',7)],
    dtype=[('name', np.str_,8), ('job', np.str_,8), ('income', np.uint32)])

result=matplotlib.mlab.rec_groupby(data, ('job',), (('income',np.mean,'avg_income'),))

产量

('Digger', 4.0)
('Planter', 2.5)
('Waterer', 3.0)

matplotlib.mlab.rec_groupby返回一个recarray：

print(result.dtype)
# [('job', '|S7'), ('avg_income', '<f8')]

---

您可能也有兴趣查看pandas ，它具有更多用于处理group-by 操作的通用工具。

Answer 3

2022 年更新：

有一个包可以很好地模拟 matlabs accumarray 的功能。 您可以通过pip install numpy_groupies安装它或在此处找到它：

https://github.com/ml31415/numpy-groupies

Answer 4

最好的灵活性和可读性是使用pandas ：

import pandas

data=np.array(
    [('Aaron','Digger',1),
     ('Bill','Planter',2),
     ('Carl','Waterer',3),
     ('Darlene','Planter',3),
     ('Earl','Digger',7)],
    dtype=[('name', np.str_,8), ('job', np.str_,8), ('income', np.uint32)])

df = pandas.DataFrame(data)
result = df.groupby('job').mean()

产生于：

         income
job
Digger      4.0
Planter     2.5
Waterer     3.0

Pandas DataFrame 是一个很棒的类，但您可以根据需要取回结果：

result.to_records()
result.to_dict()
result.to_csv()

等等...

Answer 5

使用scipy的ndimage.mean 可以获得最佳性能。 对于这个小数据集，这将比接受的答案快两倍，对于较大的输入，速度大约快 3.5 倍：

from scipy import ndimage

data=np.array(
    [('Aaron','Digger',1),
     ('Bill','Planter',2),
     ('Carl','Waterer',3),
     ('Darlene','Planter',3),
     ('Earl','Digger',7)],
    dtype=[('name', np.str_,8), ('job', np.str_,8), ('income', np.uint32)])

unique = np.unique(data['job'])
result=np.dstack([unique, ndimage.mean(data['income'], data['job'], unique)])

将屈服于：

array([[['Digger', '4.0'],
        ['Planter', '2.5'],
        ['Waterer', '3.0']]],
      dtype='|S32')

编辑：使用 bincount （更快！）

对于小示例输入，这比接受的答案快约 5 倍，如果您重复数据 100000 次，它将快约 8.5 倍：

unique, uniqueInd, uniqueCount = np.unique(data['job'], return_inverse=True, return_counts=True)
means = np.bincount(uniqueInd, data['income'])/uniqueCount
return np.dstack([unique, means])

Answer 6

http://python.net/~goodger/projects/pycon/2007/idiomatic/handout.html#dictionary-get-method

应该有助于使它更漂亮，更pythonic，更高效。 我稍后会回来检查您的进度。 也许您可以考虑到这一点来编辑该功能？ 另请参阅接下来的几节。