使用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)

这里有一个例子。请原谅，我不知道这个代码有多有效。

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 ]]

你可以使用我写的this group_by函数来解决类似的问题，如下所示：

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

print(group_by(group_id, lambda idx: a[idx].mean(axis=0))[2])
# [[1.15 1.45 2.25] [2.1  1.2  0.5 ]]

基本上，lambda先用idx=[0, 2]调用，然后用idx=[1]调用，因为组id是[0 1 0]，然后结果被堆叠在一起。
我推荐这个函数，因为它非常灵活：我也用它来处理this和other问题，这些问题需要在numpy中使用group-by、apply或transform。