Массив разделенного массива на основе условия без цикла for

Question

Допустим, у меня есть массив NumPy, который содержит точки в 2d пространстве, например:

np.array([[3, 2], [4, 4], [5, 4], [4, 2], [4, 6], [9, 5]]) 

У меня также есть массив NUMPY, который маркирует каждую точку числом, этот массив является 1D-массивом.с длиной как количество точек в массиве точек.

np.array([0, 1, 1, 0, 2, 1])

Теперь я хочу взять среднее значение каждой точки, которая имеет индекс из массива меток.Таким образом, для всех точек, имеющих метку 0, возьмите среднее значение этих точек.Мой текущий способ решения этой проблемы следующий:

return np.array([points[labels==k].mean(axis=0) for k in range(k)])

, где k - наибольшее число в массиве меток, или, как его называют, количество способов маркировки точек.

Мне бы хотелось, чтобы это можно было сделать без использования цикла for, может быть, какая-то непонятная функциональность, которую я еще не открыл?

Answer 1

Подход № 1: Мы можем использовать matrix-multiplication с некоторой помощью от braodcasting -

mask = labels == np.arange(labels.max()+1)[:,None]
out = mask.dot(points)/np.bincount(labels).astype(float)[:,None]

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

In [36]: points = np.array([[3, 2], [4, 4], [5, 4], [4, 2], [4, 6], [9, 5]]) 
    ...: labels = np.array([0, 1, 1, 0, 2, 1])

# Original soln
In [37]: L = labels.max()+1

In [38]: np.array([points[labels==k].mean(axis=0) for k in range(L)])
Out[38]: 
array([[3.5       , 2.        ],
       [6.        , 4.33333333],
       [4.        , 6.        ]])

# Proposed soln
In [39]: mask = labels == np.arange(labels.max()+1)[:,None]
    ...: out = mask.dot(points)/np.bincount(labels).astype(float)[:,None]

In [40]: out
Out[40]: 
array([[3.5       , 2.        ],
       [6.        , 4.33333333],
       [4.        , 6.        ]])

Подход № 2: С np.add.at -

sums = np.zeros((labels.max()+1,points.shape[1]),dtype=float)
np.add.at(sums,labels,points)
out = sums/np.bincount(labels).astype(float)[:,None]

Подход № 3: Если все числа из последовательностиот 0 до max-label присутствуют в labels, мы также можем использовать np.add.reduceat -

sidx = labels.argsort()
sorted_points = points[sidx]
sums = np.add.reduceat(sorted_points,np.r_[0,np.bincount(labels)[:-1].cumsum()])
out = sums/np.bincount(labels).astype(float)[:,None]