使用NumPy的argsort并采取在2D阵列

Question

目的是计算距离矩阵两组点（set1和set2）之间，使用argsort()以获得排序索引和take()提取排序后的数组。我知道我可以直接做sort()，但我需要下一步的索引。我正在使用花式索引concepts discussed here。我无法直接使用take()与获得的索引矩阵，但添加到每一行的相应数量使其工作，因为take()平整源数组使第二行元素索引+ = len（set2），第三行索引+ = 2 * LEN（设置2）等（见下文）：

dist = np.subtract.outer(set1[:,0], set2[:,0])**2 
dist += np.subtract.outer(set1[:,1], set2[:,1])**2 
dist += np.subtract.outer(set1[:,2], set2[:,2])**2 
a = np.argsort(dist, axis=1) 
a += np.array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
       [10, 10, 10, 10, 10, 10, 10, 10, 10, 10], 
       [20, 20, 20, 20, 20, 20, 20, 20, 20, 20], 
       [30, 30, 30, 30, 30, 30, 30, 30, 30, 30]]) 
s1 = np.sort(dist,axis=1) 
s2 = np.take(dist,a) 
np.nonzero((s1-s2)) == False 
#True # meaning that it works... 

的主要问题是：有没有使用take()没有这些指标相加直接的方式？

set1 = np.array([[ 250., 0., 0.], 
       [ 250., 0., 510.], 
       [-250., 0., 0.], 
       [-250., 0., 0.]]) 

set2 = np.array([[ 61.0, 243.1, 8.3], 
       [ -43.6, 246.8, 8.4], 
       [ 102.5, 228.8, 8.4], 
       [ 69.5, 240.9, 8.4], 
       [ 133.4, 212.2, 8.4], 
       [ -52.3, 245.1, 8.4], 
       [-125.8, 216.8, 8.5], 
       [-154.9, 197.1, 8.6], 
       [ 61.0, 243.1, 8.7], 
       [ -26.2, 249.3, 8.7]]) 

其他相关问题：

数据一起玩

- Euclidean distance between points in two different Numpy arrays, not within

Answer 1

我不认为存在使用np.take而不去平指数的方式。由于尺寸可能会改变，你是关闭使用np.ravel_multi_index为好，做这样的事情：

a = np.argsort(dist, axis=1) 
a = np.ravel_multi_index((np.arange(dist.shape[0])[:, None], a), dims=dist.shape) 

或者，你可以使用花哨的索引，而无需使用take：

s2 = dist[np.arange(4)[:, None], a]