Numpy：以有效的方式计算每个行子集的平均值

Question

假设有一个Hby W，2d数组x和一个长度为H的1d数组y，它指定x的每一行被分配给类y_i。比如说,

x: [[1, 2, ...],
    [3, 4, ...],
    ...         ]

y: [1, 3, 0, 2, ...]

这意味着x的第一行在类1中，第二行在类3中，依此类推

然后，我需要计算每个类的行平均数。我想出了一种方法，在课堂上循环。下面是输入和输出的示例

x = np.array([[1, 2, 3], [1, 2, 3], [3, 4, 5], [4, 5, 6]])
y = np.array([1, 0, 0, 1])

row_means = np.ndarray([k, W]) # k is the max value in y so number of classes
for i in y:
    row_means[i] = x[y == i].mean(axis=0)

row_means: array([[2. , 3. , 4. ],
                  [2.5, 3.5, 4.5]])

然而，我想不使用任何循环来计算它。有什么办法吗

Answer 1

单行解决方案将使用列表理解而不是循环，使用set(y)避免对y中的相同值重复计算：

row_means = np.array([x[y==i].mean(axis=0) for i in set(y)])

使用x和y返回以下内容：

array([[2. , 3. , 4. ],
       [2.5, 3.5, 4.5]])

如果您关心的是性能，那么此解决方案的效率大约是带循环的解决方案的1.5倍，即使使用for i in set(y)而不是for i in y。至少在我的机器上