numpy - numpy 自定义数组元素检索

Question

我有一个关于如何从 2D numpy 数组中提取某些值的问题

Foo = 
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

Bar =
array([[0, 0, 1],
       [1, 2, 3]])

我想使用 Bar 的值作为索引从 Foo 中提取元素，这样我最终得到一个Baz与Bar. 对应的i第 th 列Baz是Foo[(np.array(each j in Bar[:,i]),np.array(i,i,i,i ...))]

Baz =
array([[ 1,  2,  6],
       [ 4,  8, 12]])

我可以做几个嵌套的 for 循环，但我想知道是否有更优雅、更 numpy-ish 的方式来做到这一点。

对不起，如果这有点令人费解。让我知道是否需要进一步解释。

谢谢！

Answer 1

您可以将Bar其用作行索引，将数组[0, 1, 2]用作列索引：

# for easy copy-pasting
import numpy as np
Foo = np.array([[ 1,  2,  3], [ 4,  5,  6], [ 7,  8,  9],  [10, 11, 12]])
Bar = np.array([[0, 0, 1], [1, 2, 3]])

# now use Bar as the `i` coordinate and 0, 1, 2 as the `j` coordinate:

Foo[Bar, [0, 1, 2]]
# array([[ 1,  2,  6],
#        [ 4,  8, 12]])

# OR, to automatically generate the [0, 1, 2]

Foo[Bar, xrange(Bar.shape[1])]