Indexing NumPy 2D array with another 2D array

后端 未结 6 1721
孤独总比滥情好
孤独总比滥情好 2021-01-05 15:16

I have something like

m = array([[1, 2],
            [4, 5],
            [7, 8],
            [6, 2]])

and

select = array([         


        
6条回答
  •  悲哀的现实
    2021-01-05 15:39

    The numpy way to do this is by using np.choose or fancy indexing/take (see below):

    m = array([[1, 2],
               [4, 5],
               [7, 8],
               [6, 2]])
    select = array([0,1,0,0])
    
    result = np.choose(select, m.T)
    

    So there is no need for python loops, or anything, with all the speed advantages numpy gives you. m.T is just needed because choose is really more a choise between the two arrays np.choose(select, (m[:,0], m[:1])), but its straight forward to use it like this.


    Using fancy indexing:

    result = m[np.arange(len(select)), select]
    

    And if speed is very important np.take, which works on a 1D view (its quite a bit faster for some reason, but maybe not for these tiny arrays):

    result = m.take(select+np.arange(0, len(select) * m.shape[1], m.shape[1]))
    

提交回复
热议问题