Indexing one array by another in numpy

前端 未结 4 885
佛祖请我去吃肉
佛祖请我去吃肉 2020-11-22 05:12

Suppose I have a matrix A with some arbitrary values:

array([[ 2, 4, 5, 3],
       [ 1, 6, 8, 9],
       [ 8, 7, 0, 2]])

A

4条回答
  •  遇见更好的自我
    2020-11-22 05:35

    More recent versions have added a take_along_axis function that does the job:

    In [203]: A = np.array([[ 2, 4, 5, 3], 
         ...:        [ 1, 6, 8, 9], 
         ...:        [ 8, 7, 0, 2]])                                                
    In [204]: B = np.array([[0, 0, 1, 2], 
         ...:        [0, 3, 2, 1], 
         ...:        [3, 2, 1, 0]])                                                 
    In [205]: np.take_along_axis(A,B,1)                                             
    Out[205]: 
    array([[2, 2, 4, 5],
           [1, 9, 8, 6],
           [2, 0, 7, 8]])
    

    There's also a put_along_axis.

提交回复
热议问题