Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Welcome To Ask or Share your Answers For Others

Categories

0 votes
523 views
in Technique[技术] by (71.8m points)

python - efficient per column matrix indexing in numpy

I have two matrices of the same size, A, B. I want to use the columns of B to acsses the columns of A, on a per column basis. For example,

A = np.array([[1, 4, 7],
              [2, 5, 8],
              [3, 6, 9]])

and

B = np.array([[0, 0, 2],
              [1, 2, 1],
              [2, 1, 0]])

I want something like:

A[B] = [[1, 4, 9],  
        [2, 6, 8], 
        [3, 5, 7]]

I.e., I've used the j'th column of B as indices to the j'th column of A. Is there any effiecnt way of doing so? Thanks!

See Question&Answers more detail:os

与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
Welcome To Ask or Share your Answers For Others

1 Reply

0 votes
by (71.8m points)

You can use advanced indexing:

A[B, np.arange(A.shape[0])]

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

Or with np.take_along_axis:

np.take_along_axis(A, B, axis=0)

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

与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
OGeek|极客中国-欢迎来到极客的世界,一个免费开放的程序员编程交流平台!开放,进步,分享!让技术改变生活,让极客改变未来! Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Click Here to Ask a Question

...