从 numpy argmin 中删除轴参数，但仍然矢量化答案

【问题标题】：Removing axis argument from numpy argmin, but still vectorized从 numpy argmin 中删除轴参数，但仍然矢量化
【发布时间】：2023-03-11 13:31:01
【问题描述】：

所以我有以下代码行

np.argmin(distances, axis = 0)

这里的距离是k个质心和n个点之间的距离矩阵。所以它是一个 k x n 矩阵。因此，通过这行代码，我试图通过沿轴 0 获取 argmin 来找到每个点的最近质心。

我的目标是有一个没有轴参数的类似矢量化代码，因为它没有在我正在使用的 numpy 的分支中实现。

任何帮助都会很好:)

【问题讨论】：

分叉有np.sort(..axis)吗？
@Divakar 恐怕没有。似乎轴总是一个问题
发布的解决方案对您有用吗？

标签： python numpy vectorization

【解决方案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

【讨论】：

谢谢@Divakar，你总是那么乐于助人:) 这似乎只有一个问题。我不能使用 argsort() 函数，因为它还没有实现..我知道..很烦人。
@Shadesfear 你在什么样的原始 NumPy 分支上？ :)
我在 Bohrium bohrium.readthedocs.io 上。所以numpy代码的自动gpu加速。
@Shadesfear 是否支持数组赋值和 np.flatnonzero？
@Shadesfear 数组赋值？