利用整數數組索引根據輔助數組索引提取元素

在給定的場景中，目標是從數組A 使用第二個數組B 中指定的索引。更直接的方法是使用NumPy 的整數陣列索引，而不是依賴np.take 或np.choose：

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

以下是此程式碼如何實現所需結果的：

• np.arange(A.shape[0])建立一個陣列從0 到(A.shape[0] - 1) 的索引，表示A.
• B.ravel() 的行，將B 展平為一維數組，確保其具有相同數量的元素作為A 中的行。
• 索引操作 A[...] 使用 B 中的行索引和列索引檢索 A 中的元素。

當 B 時，此方法特別有用是一維數組或列索引列表。透過跳過展平操作，程式碼變得更簡單：

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]