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) 向量。）

在 numpy 中是否有任何方法可以做到这一点，或者我是否需要自己实现二进制搜索（仍然是一个选项，因为它即使在 vanilla 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]]

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

Numpy searchsorted on array of vectors in one dimension

python

arrays

numpy

vector

binary-search