numpy自定义数组元素检索

时间:2012-11-20 00:21:06

标签: numpy

我有一个关于如何从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形状相同的2D矩阵/数组Bari对应的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的方法来做到这一点。

很抱歉,如果这有点令人费解。如果我需要进一步解释,请告诉我。

谢谢!

1 个答案:

答案 0 :(得分:2)

您可以使用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])]