Python: 查找索引

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是一维或者二维数组怎么办？说，

b=[8.1,3.1,9.1]

我想得到 (2,1),(0,2),(2,2)

一般来说，对于 b 的每个值，我希望 a 中只有一个匹配项。我可以避免 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 从 1D 扩展 b 到 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: 查找索引

Python: finding indices

python

arrays

numpy

vectorization

find