整数配列インデックスを利用してセカンダリ配列インデックスに基づいて要素を抽出する

指定されたシナリオの目標は、配列から特定の要素を取得することです。 2 番目の配列 B で指定されたインデックスを使用して配列 A を作成します。 np.take または np.choose に依存する代わりに、より簡単なアプローチは、NumPy の整数配列インデックスを使用することです:

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

このコードが望ましい結果を達成する方法は次のとおりです。

• np.arange(A.shape[0]) は次の配列を作成します。 0 から (A.shape[0] - 1) までのインデックスは A の行を表します。
• B.ravel() は B を 1 次元配列に平坦化し、要素数が同じになるようにします。
• インデックス操作 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]