numpy argmax与groupby

Question

假设我有标签将数据聚集到几个组中。现在我想找出每个簇的原始数组的最大w.r.t索引。例如：

import numpy as np
labels = np.array([1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 0])
y = np.random.randint(0, 9, len(labels)) #array([6, 7, 5, 4, 2, 8, 4, 4, 5, 6, 4])

我想要1，5，因为对于簇1，在索引1处，最大值是7，对于簇0，在索引5处，最大值是8，如果没有for循环，可以得到吗？

我的天真的解决方案供参考：

out = []
for i in [0, 1]:
    temp = y.copy()
    temp[labels == i] = -1
    out.append(np.argmax(temp))

Answer 1

我相信这应该是可行的：

labels = np.array([1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 0])
y = np.array([6, 7, 5, 4, 2, 8, 4, 4, 5, 6, 4])
arr = np.unique(labels)

result = np.ma.masked_where(labels[:,None] == arr, np.tile(y,(arr.shape[0],1)).T).argmax(axis=0)

其主要思想是，我们通过添加新的维度来强制一维数组进行广播。