我有一千个128维特征的数据集，其形状为（1000,128）。
我想找到形状为（128,1）的128维特征的排序最近的邻居。
通过数据集（1000,128）与要素（128,1）之间的矩阵乘法计算的距离，将得出形状相似的数组（1000,1）：
这是通过以下方式完成的：
# features.shape=(1000,128) ; feature.shape=(128,1) ; similarities.shape=(1000,1) similarities = features.dot(feature)
在计算了距离（相似度）之后，我使用以下代码找到了最近的邻居：
# The n Nearest Neighbors Indexes (But Not Sorted) nearest_neighbours_indexes_unsorted = np.argpartition(similarities, kth=-n)[-n:] # The n Nearest Neighbors (But Not Sorted) nearest_neighbours_similarities_unsorted = similarities[nearest_neighbours_indexes_unsorted] # The Indexes of n Nearest Neighbors Sorted nearest_neighbours_indexes_sorted = np.flip(nearest_neighbours_indexes_unsorted[np.argsort(nearest_neighbours_similarities_unsorted)], axis=0)
这段代码可快速处理数百万个数据（我想知道是否有人提示使其更快），但我希望能够一口气找到多个功能的最近邻居：
一种方法是为一个循环中的每个功能计算上述代码（这很慢），另一种方法是更改代码以适应多维索引，这就是我遇到的问题：我不知道如何编写上面的代码针对形状为（128，n）而非（128,1）的特征。
这是一个辅助函数，n-largest可使用np.argpartition和np.take_along_axis-从通用ndarray沿通用轴选择顶部索引。
n-largest
np.argpartition
np.take_along_axis
def take_largest_indices_along_axis(ar, n, axis): s = ar.ndim*[slice(None,None,None)] s[axis] = slice(-n,None,None) idx = np.argpartition(ar, kth=-n, axis=axis)[tuple(s)] sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis) return np.flip(np.take_along_axis(idx, sidx, axis=axis),axis=axis)
扩展它以获得n个最小的索引-
def take_smallest_indices_along_axis(ar, n, axis): s = ar.ndim*[slice(None,None,None)] s[axis] = slice(None,n,None) idx = np.argpartition(ar, kth=n, axis=axis)[tuple(s)] sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis) return np.take_along_axis(idx, sidx, axis=axis)
并扩展这些n元素以选择最大或最小的元素本身，将使用下面np.take_along_axis列出的简单用法-
n
def take_largest_along_axis(ar, n, axis): idx = take_largest_indices_along_axis(ar, n, axis) return np.take_along_axis(ar, idx, axis=axis) def take_smallest_along_axis(ar, n, axis): idx = take_smallest_indices_along_axis(ar, n, axis) return np.take_along_axis(ar, idx, axis=axis)
样品运行
# Sample setup In [200]: np.random.seed(0) ...: ar = np.random.randint(0,99,(5,5)) In [201]: ar Out[201]: array([[44, 47, 64, 67, 67], [ 9, 83, 21, 36, 87], [70, 88, 88, 12, 58], [65, 39, 87, 46, 88], [81, 37, 25, 77, 72]])
n沿轴取最大索引，元素-
In [202]: take_largest_indices_along_axis(ar, n=2, axis=0) Out[202]: array([[4, 2, 2, 4, 3], [2, 1, 3, 0, 1]]) In [203]: take_largest_indices_along_axis(ar, n=2, axis=1) Out[203]: array([[4, 3], [4, 1], [2, 1], [4, 2], [0, 3]]) In [251]: take_largest_along_axis(ar, n=2, axis=0) Out[251]: array([[81, 88, 88, 77, 88], [70, 83, 87, 67, 87]]) In [252]: take_largest_along_axis(ar, n=2, axis=1) Out[252]: array([[67, 67], [87, 83], [88, 88], [88, 87], [81, 77]])
n沿轴取最小的索引，元素-
In [232]: take_smallest_indices_along_axis(ar, n=2, axis=0) Out[232]: array([[1, 4, 1, 2, 2], [0, 3, 4, 1, 0]]) In [233]: take_smallest_indices_along_axis(ar, n=2, axis=1) Out[233]: array([[0, 1], [0, 2], [3, 4], [1, 3], [2, 1]]) In [253]: take_smallest_along_axis(ar, n=2, axis=0) Out[253]: array([[ 9, 37, 21, 12, 58], [44, 39, 25, 36, 67]]) In [254]: take_smallest_along_axis(ar, n=2, axis=1) Out[254]: array([[44, 47], [ 9, 21], [12, 58], [39, 46], [25, 37]])
对于我们的情况，假设输入similarities的形状是，并且(1000,128)表示1000个数据点和128个要素，并且我们要n=10为每个数据点寻找最大的说要素，那么它将是-
similarities
(1000,128)
n=10
take_largest_indices_along_axis(similarities, n=10, axis=1) # indices take_largest_along_axis(similarities, n=10, axis=1) # elements
最终的索引/值数组将为shape (1000, n)。
(1000, n)
以给定的数据集形状运行样本-
In [257]: np.random.seed(0) ...: similarities = np.random.randint(0,99,(1000,128)) In [263]: take_largest_indices_along_axis(similarities, n=10, axis=1).shape Out[263]: (1000, 10) In [264]: take_largest_along_axis(similarities, n=10, axis=1).shape Out[264]: (1000, 10)
相反，如果您希望获得n每个功能的最大数据点，则最终的index / values数组将为shape (n, 128)，请使用axis=0。
(n, 128)
axis=0
