Python:寻找索引

时间:2016-01-08 23:03:08

标签: python arrays numpy find vectorization

如果我有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循环吗?

2 个答案:

答案 0 :(得分:3)

使用列表理解:

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

答案 1 :(得分:2)

NumPy's broadcasting -

的矢量化方法
np.argwhere((np.abs(a - b[:,None,None])<0.5))[:,1:]

说明 -

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

    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]])