Numpy searchsorted 在一维向量数组上答案

【问题标题】：Numpy searchsorted on array of vectors in one dimensionNumpy searchsorted 在一维向量数组上
【发布时间】：2019-08-07 17:31:37
【问题描述】：

我已经在谷歌上搜索了一段时间，尽管我认为这是一个常见问题，但我认为在 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 中执行此操作，或者我是否需要自己实现二进制搜索（仍然是一个选项，因为即使在普通 Python 中它也非常快）。

例如对于值x= 2.5，我希望索引为 2。对于 x=2.4，我仍然希望为 2，对于 x=2.6，我希望为 3。索引或向量本身（例如 ( 3、1、2.5))。

【问题讨论】：

有任何理由使用元组数组吗？如果您要使用形状 (4,3) 的 numpy 数组，如果您只是对某个组件感兴趣，则可以将其切片。
@nahtanoj 我只使用元组，因为如果我将数组放入其中，则排序不起作用。如果你有一个两者都可以工作的例子，我会很高兴的。

标签： python arrays numpy vector binary-search

【解决方案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]]

【讨论】：