在Python中按特征求和数组

时间:2011-09-14 13:15:01

标签: python arrays numpy sum

我想知道通过给定特征对数组元素求和的最有效方法是什么。例如,我有1000个数据,我正在寻找的是给定年份疾病的每次抽奖(列)的总和(即抽奖是按性别,年份,疾病,我想要的每年的性别和疾病的总和)。

import numpy as np
year = np.repeat((1980, 1990 , 2000, 2010), 10)
sex = np.array(['male', 'female']*20)
disease = np.repeat(('d1', 'd2', 'd3', 'd4', 'd5', 'd6', 'd7', 'd8'), 5)
draws = np.random.normal(0, 1, size=(sex.shape[0], 1000))

关于如何获得一个形状(20,1000)的阵列的任何想法,这个阵列具有针对给定年份疾病的两种性别的平局总和?我还需要能够在数据不完全正方形的情况下做到这一点(有疾病年只有一种性别)。

2 个答案:

答案 0 :(得分:1)

import numpy as np
import itertools   
import csv

year = np.repeat((1980, 1990 , 2000, 2010), 10)
sex = np.array(['male', 'female']*20)
disease = np.repeat(('d1', 'd2', 'd3', 'd4', 'd5', 'd6', 'd7', 'd8'), 5)
draws = np.random.normal(0, 1, size=(sex.shape[0], 1000))

years=np.unique(year)
diseases=np.unique(disease)

draw_sums = dict(((y,d), draws[(year==y)&(disease==d)].sum(axis=0)) 
                  for y,d in itertools.product(years,diseases))

这导致dict将每个(年份,疾病)与相应的抽奖总和相关联。要将draw_sums写入csv,您可以执行以下操作:

with open('/tmp/test.csv','w') as f:
    writer=csv.writer(f)
    writer.writerow(['year', 'date']+['draw{i}'.format(i=i) for i in range(1,1001)])
    for yeardate,draws in sorted(draw_sums.items()):
        writer.writerow(list(yeardate)+draws.tolist())

答案 1 :(得分:0)

这是一个典型的分组问题,可以使用numpy_indexed包以完全向量化的方式有效地解决(免责声明:我是其作者)

keys, values = npi.group_by((year, disease)).sum(draws)
for key, value in zip(zip(*keys), values):
    print(key, value.shape)