如果我有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循环吗?
答案 0 :(得分:3)
使用列表理解:
[nonzero(abs(x-a)<0.5) for x in b]
答案 1 :(得分:2)
np.argwhere((np.abs(a - b[:,None,None])<0.5))[:,1:]
说明 -
None/np.newaxis
将b
从1D扩展为3D案例,并使元素沿第一轴保持。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]])