numpy group by index

Question

a = np.array([[1.2, 1.4, 2.2], [2.1, 1.2, 0.5],[1.1, 1.5, 2.3]])

group_id= np.array([0,1,0])

group id定义数组所属的组，对于每个组，计算数组的平均值。

预期结果应为：

[[1.15,1.45,2.25], [2.1,1.2,0.5]]

我怎样才能有效地做到这一点？

Answer 1

这是一个例子。请原谅我不确定这段代码的效率如何。

import numpy as np
a = np.array([[1.2, 1.4, 2.2], [2.1, 1.2, 0.5],[1.1, 1.5, 2.3]])

group_id= np.array([0,1,0])
unique_id = np.unique(group_id)

# prepare a variable
num_elements = a.shape[1]
num_id = len(unique_id)
results = np.zeros((num_id, num_elements))

# iterate over a set of unique id
for id, result in zip(unique_id, results):
    rows = np.where(group_id == id)
    result += a[rows].mean(axis=0)

print(results)

你的收益率：

[[1.15 1.45 2.25]
 [2.1  1.2  0.5 ]]

Answer 2

使用1-hot编码对代码进行矢量化处理。

import numpy as np

a = np.array([[1.2, 1.4, 2.2], [2.1, 1.2, 0.5],[1.1, 1.5, 2.3]])
group_id= np.array([0,1,0])

def onehot(arr):
    ''' Returns a 1-hot encoding array of arr
    With no zero groups, returns sorted labels
    '''
    arr_local=np.copy(arr)
    unique_arr=np.unique(arr_local)
    one_hot=(arr_local.reshape(-1,1)==unique_arr.reshape(1,-1)).astype(np.int)
    return one_hot,unique_arr

one_hot,_=onehot(group_id)
one_hot.T@a/one_hot.sum(axis=0).reshape(-1,1)