Суммирование массивов по характеристикам в Python

Question

Мне интересно, как наиболее эффективно суммировать элементы массива по заданным характеристикам. Например, у меня есть 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), который имеет сумму ничьей между обоими полами для данного года болезни? Я также должен быть в состоянии сделать это в ситуациях, когда данные не являются совершенно квадратными (есть годы заболевания, у которых есть только один пол).

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

Это приводит к диктату, связывающему каждого (год, болезнь) с соответствующей суммой розыгрышей. Чтобы записать 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())

Answer 2

Это типичная групповая проблема, которая может быть эффективно решена полностью векторизованным способом с использованием пакета numpy_indexed (отказ от ответственности: я являюсь его автором)

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