Find matching rows in 2 dimensional numpy array

前端 未结 4 1158
心在旅途
心在旅途 2020-11-30 01:38

I would like to get the index of a 2 dimensional Numpy array that matches a row. For example, my array is this:

vals = np.array([[0, 0],
                 [1         


        
4条回答
  •  Happy的楠姐
    2020-11-30 02:20

    In [5]: np.where((vals[:,0] == 0) & (vals[:,1]==1))[0]
    Out[5]: array([ 3, 15])
    

    I'm not sure why, but this is significantly faster than
    np.where((vals == (0, 1)).all(axis=1)):

    In [34]: vals2 = np.tile(vals, (1000,1))
    
    In [35]: %timeit np.where((vals2 == (0, 1)).all(axis=1))[0]
    1000 loops, best of 3: 808 µs per loop
    
    In [36]: %timeit np.where((vals2[:,0] == 0) & (vals2[:,1]==1))[0]
    10000 loops, best of 3: 152 µs per loop
    

提交回复
热议问题