Question

如果我有a和b：

a=[[1,2,3],
   [4,5,6],
   [7,8,9]]

b=8.1

我希望在a中找到值b的索引，我可以这样做：

nonzero(abs(a-b)<0.5)

获取（2,1）作为索引，但如果b是1d或2d数组，我该怎么办？说，

b=[8.1,3.1,9.1]

我希望得到（2,1），（0,2），（2,2）

一般来说，我希望b的每个值只有一个匹配。我可以避免使用for循环吗？

Answer 1

使用列表理解：

[nonzero(abs(x-a)<0.5) for x in b]

Answer 2

NumPy's broadcasting -

的矢量化方法

np.argwhere((np.abs(a - b[:,None,None])<0.5))[:,1:]

说明 -

使用None/np.newaxis将b从1D扩展为3D案例，并使元素沿第一轴保持。
使用2D数组a执行绝对减法，从而引入广播并导致元素减法的3D数组。
与0.5的阈值进行比较，得到与最后两个轴匹配的索引，并按第一个轴np.argwhere(...)[:,1:]排序。

示例运行 -

In [71]: a
Out[71]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [72]: b
Out[72]: array([ 8.1,  3.1,  9.1,  0.7])

In [73]: np.argwhere((np.abs(a - b[:,None,None])<0.5))[:,1:]
Out[73]: 
array([[2, 1],
       [0, 2],
       [2, 2],
       [0, 0]])

Python：寻找索引

2 个答案: