정수 배열 인덱싱을 활용하여 보조 배열 인덱스를 기반으로 요소 추출

주어진 시나리오에서 목표는 배열에서 특정 요소를 검색하는 것입니다. 두 번째 배열 B에 지정된 인덱스를 사용하여 배열 A. np.take 또는 np.choose에 의존하는 대신 더 간단한 접근 방식은 NumPy의 정수 배열 인덱싱을 사용하는 것입니다:

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

이 코드가 원하는 결과를 얻는 방법은 다음과 같습니다.

• np.arange(A.shape[0])는 다음과 같은 배열을 생성합니다. A.
• B.ravel() 행을 나타내는 0부터 (A.shape[0] - 1)까지의 인덱스는 B를 1차원 배열로 병합하여 동일한 수의 요소를 갖도록 보장합니다. A의 행으로.
• 인덱싱 작업 A[...]는 B의 행 인덱스와 열 인덱스를 사용하여 A에서 요소를 검색합니다.

이 접근 방식은 B에서 특히 유용합니다. 1D 배열이거나 열 인덱스 목록입니다. 병합 작업을 건너뛰면 코드가 더욱 단순해집니다.

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

예:

A = np.array([[0, 1], [2, 3], [4, 5]])
B = np.array([1, 0, 1])

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

print(result)
# Output: [1, 2, 5]