numpy数组中值数组的位置

Question

这里有一个小代码来说明问题。

A = array([[1,2], [1,0], [5,3]])
f_of_A = f(A)   # this is precomputed and expensive


values = array([[1,2], [1,0]])


# location of values in A
# if I just had 1d values I could use numpy.in1d here
indices = array([0, 1])


# example of operation type I need (recalculating f_of_A as needed is not an option)
f_of_A[ indices ]

所以，基本上我认为我需要一些等效于 in1d 的更高维度。这样的事情存在吗？还是有其他方法？

看起来还有一个 searchsorted() 函数，但这似乎也适用于一维数组。在此示例中，我使用了 2d 点，但任何解决方案也需要适用于 3d 点。

Answer 1

好的，这就是我想出的。

要找到一个多维索引的值，比如说ii = np.array([1,2])，我们可以这样做：

n.where((A == ii).all(axis=1))[0]

让我们分解一下，我们有A == ii，它将针对A 的每一行与ii 进行元素比较。我们希望整行都是真的，所以我们添加.all(axis=1) 来折叠它们。为了找到这些索引发生的位置，我们将其插入np.where 并获取元组的第一个值。

现在，我还没有一种快速的方法来使用多个索引（尽管我感觉有一个）。但是，这将完成工作：

np.hstack([np.where((A == values[i]).all(axis=1))[0] for i in xrange(len(values))])

这基本上只是针对values 的每个值调用上述方法，并将结果连接起来。

更新：

这里是多维案例（一次完成，应该相当快）：

np.where((np.expand_dims(A, -1) == values.T).all(axis=1).any(axis=1))[0]

Answer 2

您可以在原始数组的视图上使用np.in1d，其中所有坐标都折叠到一个dtype np.void的变量中：

import numpy as np

A = np.array([[1,2], [1,0], [5,3]])
values = np.array([[1,2], [1,0]])

# Make sure both arrays are contiguous and have common dtype
common_dtype = np.common_type(A, values)
a = np.ascontiguousarray(A, dtype=common_dtype)
vals = np.ascontiguousarray(values, dtype=common_dtype)

a_view = A.view((np.void, A.dtype.itemsize*A.shape[1])).ravel()
values_view = values.view((np.void,
                           values.dtype.itemsize*values.shape[1])).ravel()

现在a_view 和values_view 的每个项目都是一个点的所有坐标打包在一起，所以你可以做任何你会使用的一维魔法。我不知道如何使用np.in1d 来查找索引，所以我会选择np.searchsorted 路线：

sort_idx = np.argsort(a_view)
locations = np.searchsorted(a_view, values_view, sorter=sort_idx)
locations = sort_idx[locations]

>>> locations
array([0, 1], dtype=int64)