在一维向量数组上进行的Numpy搜索排序

Question

我已经使用谷歌搜索了一段时间，尽管我认为这是一个普遍的问题，但我看不出在SO上有任何解决方案。

假设我有一个3D向量数组（x，y，z），如下所示：

import numpy as np

arr = np.array(
    [(1, 2, 3), (3, 1, 2.5), (5, 3, 1), (0, -1, 2)],
    dtype=[('x', np.float), ('y', np.float), ('z', np.float)]
)
print(np.sort(arr, order='z'))

打印：

[(5.,  3., 1. ) (0., -1., 2. ) (3.,  1., 2.5) (1.,  2., 3. )]

我现在想仅按维度“ z”搜索此排序数组。 二进制搜索将非常有效。 但是searchsorted仅适用于一维数组。 而且没有lambda可以应用于每个值（基本上np.dot带有(0, 0, 1) np.dot (0, 0, 1)向量的np.dot 。）

是否有任何方法可以在numpy中执行此操作，或者我需要自己实现二进制搜索（仍然是一个选项，因为即使在原始Python中它也非常快）。

例如，对于值x= 2.5我期望索引2。对于x=2.4我仍然期望2，对于x=2.6我期望3。索引或向量本身（例如（3，1 ，2.5））。

Answer 1

无需在数组中使用元组，就可以使用切片：

import numpy as np

arr = np.random.rand(10,3)
print(arr)

sort_indices = np.argsort(arr[:,2])
arr_sorted = arr[sort_indices]
print(arr_sorted)


# run search sorted
search_result = np.searchsorted(arr_sorted[:,2],arr[5,2])
>>> 2

输出：

unsorted:
[[0.71815835 0.89099775 0.51398111]
 [0.56393906 0.26684628 0.33065586]
 [0.38920018 0.0485013  0.70958811]
 [0.3771277  0.95567051 0.18514701]
 [0.59715961 0.19092995 0.09340359]
 [0.09575273 0.56697649 0.10120321]
 [0.63226061 0.95258914 0.59669295]
 [0.1714133  0.7406211  0.23079041]
 [0.33512727 0.23244954 0.08735154]
 [0.50582011 0.97186928 0.15525005]]

sorted:
[[0.33512727 0.23244954 0.08735154]
 [0.59715961 0.19092995 0.09340359]
 [0.09575273 0.56697649 0.10120321]
 [0.50582011 0.97186928 0.15525005]
 [0.3771277  0.95567051 0.18514701]
 [0.1714133  0.7406211  0.23079041]
 [0.56393906 0.26684628 0.33065586]
 [0.71815835 0.89099775 0.51398111]
 [0.63226061 0.95258914 0.59669295]
 [0.38920018 0.0485013  0.70958811]]