공상 색인 생성 numpy 행렬 : 행당 하나의 요소

Question

나는 shape (m, n)의 2d numpy 배열, 행렬을 가지고있다. 더불어

matrix = np.arange(5*3).reshape((5, 3)) 

I 모양 (m)의 정수의 인덱싱 어레이 IDX을 가지고 내 실제 유스 케이스는 m ~ 1E5 및 N ~ 100를 가지고 있지만, 단순한 최소 예를 갖는 위해서 [0, n] 사이의 각 항목. 이 배열은 행의 각 행에서 선택할 열을 지정합니다.

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

그래서, 따라서 4 행에서 1 행 및 열 1 행 2 열 1 행 1 열 2에서 행 0, 열 0에서 2 열을 선택하려면 최종 답변을해야합니다 노력하고 있어요 수 :

correct_result = np.array((2, 3, 8, 10, 13)) 

나는 직관적으로 다음을 시도했지만 부정확 한 :

incorrect_result = matrix[:, idx] 

은 무엇 위의 구문이하는 일은 행으로 멋진 색인 배열, 행으로 IDX을 적용하고, 다른 결과 형상의 매트릭스 (m, n), 그건 내가 원하는 바가 아니야.

이 유형의 고급 색인 생성에 올바른 구문은 무엇입니까?

Answer 1

correct_result = matrix[np.arange(m), idx] 

고급 인덱싱 식 matrix[I, J]output[n] == matrix[I[n], J[n]]는 그러한 출력을 제공한다. 우리가 output[n] == matrix[n, idx[n]]을 원하는 경우에

는, 우리는 I[n] == n 및 J[n] == idx[n] 필요하므로 idx가 될 우리는 np.arange(m) 및 J을 할 I이 필요합니다.