[解決済み] NumPyはインデックスのリストを使って、行ごとに特定の列インデックスを選択します。

質問

NumPyの行列の行ごとの特定の列を選択するのに苦労しています。

私が以下のような行列を持っているとします。 X :

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

また list と呼んでいる、各行ごとのカラムインデックスの Y :

[1, 0, 2]

値を取得したい

[2]
[4]
[9]

の代わりに list というインデックスを持つ Y と同じ形状の行列を作ることもできます。 X ここで，各列は bool / int のように、0-1の値で、これが必須のカラムであるかどうかを示します。

[0, 1, 0]
[1, 0, 0]
[0, 0, 1]

配列に対して反復処理を行い、必要な列の値を選択することで実現できることは知っています。しかし、これは大きなデータの配列に対して頻繁に実行されるため、できる限り高速に実行する必要があります。

したがって、私はより良いソリューションがあるかどうか疑問に思っていました。

どのように解決するのですか？

boolean配列があれば、それを元に以下のように直接選択することができます。

>>> a = np.array([True, True, True, False, False])
>>> b = np.array([1,2,3,4,5])
>>> b[a]
array([1, 2, 3])

最初の例に沿って、次のようにすることができます。

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> b = np.array([[False,True,False],[True,False,False],[False,False,True]])
>>> a[b]
array([2, 4, 9])

また arange を追加して、それを直接選択することもできます。

>>> a = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> a[np.arange(len(a)), [1,0,2]]
array([2, 4, 9])

お役に立てれば幸いです。また何か質問があれば教えてください。