按元素上的条件值对 Numpy 数组元素进行排序

Question

我需要通过增加与另一个点的距离来对点的 numpy 数组进行排序。

import numpy as np
def dist(i,j,ip,jp): 
return np.sqrt((i-ip)**2+(j-jp)**2)
arr = np.array([[0,0],[1,2],[4,1]])

我想做的是在 arr 中的固定点 [i，j]=[1,1] 和每个有序对 [ip，jp] 之间使用函数 dist(1,1，ip，jp) 返回 arr，每个元素从最低到最高接近 [i，j] 排序。有人对此有快速解决方案吗？

我想要的输出是 new_arr = np.array([[1,2]，[0,0]，[4,1]])

我有一些想法，但它们似乎效率很低。

谢谢！

Answer 1

似乎有两种方法可以做到这一点：

1. 将整个 numpy 数组转换为 Python 列表，并使用 Python 的排序方法和键函数对其进行排序。

   l = list(arr)
   l.sort(key=lambda coord: dist(1, 1, coord[0], coord[1]))
   arr = np.array(l)
   

2. 通过在原始数组上映射dist()来创建第二个 numpy 数组，使用.argsort()获取排序顺序，然后将其应用于原始数组。

   arr2 = np.vectorize(lambda coord: dist(1, 1, coord[0], coord[1]))(arr)
   arr3 = np.argsort(arr2)
   arr = np.array(arr)[arr3]
   

Answer 2

实际上，您可以使用numpy.lexsort()，使用带有dist()作为键的arr映射，然后将生成的掩码应用于arr本身，从而使@user3030010的第二个答案更有效率

。

import numpy as np
my_key = np.vectorize(lambda coord: dist(1, 1, coord[0], coord[1]))(arr)
inds = np.lexsort(keys = [my_key])
arr = arr[inds]

这确实是一个很小的改进，但如果您随后添加更多键进行排序，该方法特别有用。

Answer 3

我现在意识到这是一个流行的问题，所以多年后，这是我自己的答案，它使用了scipy.spatial极其强大的功能。在这里，scipy.spatial.cdist用于进行距离计算。这是闪电般的快速和pythonic，没有任何"转换为列表并转换回来"的黑客：

from scipy.spatial.distance import cdist
import numpy as np
# EXAMPLE DATA
arr = 20*np.random.random(size=(5000000,2))-10 # testing data
pt = np.array([[1,1]]) # the point to eval proximity to
# THE SOLUTION
out = arr[np.argsort(cdist(arr,pt).squeeze())]

在这里，cdist获取一个距离数组，squeeze杀死此数组中的额外维度，argsort按距离将索引排序到距离中，arr[...]按这些索引对arr进行排序。