python - 基于其他数组访问 Numpy 数组

Question

我有一个大小为 [2, 10] 的 numpy 数组 A。我有另一个长度为 10 的数组 B，其值范围在 0 到 1 之间。我现在想要访问 A，为每列获取 B 指定的行，这意味着:输出的长度应为 10，如果 B[i] == 0，则结​​果数组中的第 i 个条目为 A[0, i]，否则为 A[1, i]。我如何正确/最优雅地对其进行索引？谢谢。

Answer 1

您可以使用numpy.choose为此:

>>> a
array([[ -3,   2,   5,   0,  -2,  12,   5,  -1, -12,   5],
       [ -8, -12,  15,   6,   2,  -3,   4, -17,   0,   7]])
>>> b
array([1, 0, 0, 1, 1, 1, 1, 1, 0, 1])
>>> np.choose(b, a)
array([ -8,   2,   5,   6,   2,  -3,   4, -17, -12,   7])

这里第一个参数 b 是包含选项值的数组，下一个参数是每个选项值的二维数组。

如您所见，对于索引i，其中b[i] == 0，则第一行的相应值a[0] [i]，对于 b[i] == 1，我们得到 a[1][i]。这可以很容易地推广到更多值。

或者我们可以使用numpy.where并传递子数组:

>>> np.where(b, a[1], a[0])
array([ -8,   2,   5,   6,   2,  -3,   4, -17, -12,   7])