Разделите значения NumPy по индексам, что приведет к неправильной форме

Question

Вот чего я хочу достичь с помощью NumPy и понятия не имею, как.Чтобы быть ясным, я хотел бы сделать это максимально кратко.

# shape (5, 2)
data = np.array([
    [10, 20]
    [30, 50]
    [10, 10]
    [5, 13]
    [7, 7]
])

# shape (5,)
target = np.array([0, 2, 1, 0, 2])

# how to achieve this in a more numpy way
# shape(3, 2)
result = np.zeros((target.max() + 1, data.shape[1]))
for i in range(result.shape[0])):
    result[i] = data[np.where(target == i)].mean(axis=0)

Я знаю, что это может быть один вкладыш:

result = np.array([data[np.where(target == i)].mean(axis=0)
                   for i in range(target.max() + 1)])

Спасибо

РЕДАКТИРОВАТЬ: Там был опечатка в цикле.

Answer 1

Пакет numpy_indexed (заявление об отказе: я его автор) был создан для эффективного и краткого решения таких проблем, как:

import numpy_indexed as npi
unique, result = npi.group_by(target).mean(data)

Answer 2

Подход № 1: Мы можем использовать np.add.at -

def binmean_addat(data, target):
    result = np.zeros((target.max() + 1, data.shape[1]))
    np.add.at(result, target,data)
    grouped_count = np.bincount(target)
    out = result/np.bincount(target)[:,None]
    out[grouped_count==0] = 0
    return out

Подход № 2: Мы также можем использовать matrix-multiplication -

def binmean_dot(data, target):
    grouped_sum = (target == np.arange(target.max() + 1)[:,None]).dot(data)
    grouped_count = np.bincount(target)
    out = np.true_divide(grouped_sum,grouped_count[:,None])
    out[grouped_count==0] = 0
    return out

Подход № 3: С np.add.reduceat -

def binmean_reduceat(data, target):
    sidx = target.argsort()
    grouped_count = np.bincount(target)
    grouped_sum = np.add.reduceat(data[sidx],np.r_[0,grouped_count[:-1].cumsum()])
    out = np.true_divide(grouped_sum,grouped_count[:,None])
    out[grouped_count==0] = 0
    return out

Пробный прогон -

In [45]: data
Out[45]: 
array([[10, 20],
       [30, 50],
       [10, 10],
       [ 5, 13],
       [ 7,  7]])

In [46]: target
Out[46]: array([0, 2, 1, 0, 2])

In [47]: binmean_addat(data, target)
Out[47]: 
array([[ 7.5, 16.5],
       [10. , 10. ],
       [18.5, 28.5]])

---

Для детей, играющих code-golf, вот два (один из них - измененная версия из опубликованного вопроса, а другой - из одного из предыдущих в этом посте) -

Приложение № 1 -

(data.T.dot(target[:,None]==range(target.max()+1))/(np.bincount(target)+0.0)).T

Приложение № 2 -

np.array([data[target==t].mean(0) for t in range(target.max()+1)])