Удаление аргумента оси из numpy argmin, но все еще векторизация

Question

Итак, у меня есть следующие строки кода

np.argmin(distances, axis = 0)

Здесь расстояния - это матрица расстояний между k центроидами и n точками. так что это матрица k x n. Итак, с помощью этой строки кода я пытаюсь найти ближайший центроид для каждой точки, беря argmin вдоль оси 0.

Моя цель - получить похожий векторизованный код без аргумента оси, так как он не реализован в развилке numpy, которую я использую.

Любая помощь была бы хороша :)

Answer 1

Вот векторизованный -

def partial_argsort(a):
    idar = np.zeros(a.max()+1,dtype=int)
    idar[a] = np.arange(len(a))
    return idar[np.sort(a)]

def argmin_0(a):
    # Define a scaling array to scale each col such that each col is 
    # offsetted against its previous one  
    s = (a.max()+1)*np.arange(a.shape[1])

    # Scale each col, flatten with col-major order. Find global partial-argsort. 
    # With the offsetting, those argsort indices would be limited to per-col
    # Subtract each group of ncols elements based on the offsetting.
    m,n = a.shape
    a1D = (a+s).T.ravel()
    return partial_argsort(a1D)[::m]-m*np.arange(n)

Пример прогона для проверки -

In [442]: np.random.seed(0)
     ...: a = np.random.randint(11,9999,(1000,1000))
     ...: idx0 = argmin_0(a)
     ...: idx1 = a.argmin(0)
     ...: r = np.arange(len(idx0))
     ...: print (a[idx0,r] == a[idx1,r]).all()
True