Question

我们说我有一个二维numpy数组，

import numpy as np
a = np.random.rand(5,5)

和一维索引列表

b = np.array([0,4,2,3,3])

我想要做的是生成a的值数组，如下所示：

[a[b[0],0], a[b[1],1], a[b[2],[2]], a[b[3],3], a[b[4],4]]

现在，我知道我可以使用python＆＃39;

来实现这一目标

[a[b[i], i] for i in range(5)]

但这对我来说有两个问题：

比纯粹numpy＆＃34;中执行的操作要慢得多。我认为可以通过numpy.choose的一些花哨使用来避免，但这仍然让我遇到第二个问题
在我的程序的后续部分中，我实际上想要在适当的位置更改a的值，即我生成大小为x的向量（数组）5，然后设置对于a[b[i],i]的所有五个falues，x到i。

目前，我使用for循环执行此操作：

x = np.random.rand(5)
for i in range(5):
    a[b[i], i] = x[i]

但我觉得这可以更快地完成。任何帮助将不胜感激。

谢谢。

Answer 1

获得：

>>> a = np.arange(25).reshape((5,5))
>>> a
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])
>>> rows = b = np.array([0,4,2,3,3])
>>> cols = np.arange(len(b))
>>> [a[b[i], i] for i in range(5)]
[0, 21, 12, 18, 19]
>>> a[rows, cols]
array([ 0, 21, 12, 18, 19])

环境：

>>> a[rows, cols] = 69
>>> a
array([[69,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 69, 13, 14],
       [15, 16, 17, 69, 69],
       [20, 69, 22, 23, 24]])
>>> a[rows, cols] = np.array([-111, -111, -222, -333, -666])
>>> a
array([[-111,    1,    2,    3,    4],
       [   5,    6,    7,    8,    9],
       [  10,   11, -222,   13,   14],
       [  15,   16,   17, -333, -666],
       [  20, -111,   22,   23,   24]])

Answer 2

>>> import numpy as np
>>> a = np.random.rand(5,5)
>>> b = np.array([0,4,2,3,3])
>>> b
array([[ 0.4564633 ,  0.72926453,  0.74926462,  0.41773619,  0.46413324],
       [ 0.17443626,  0.82383088,  0.39718456,  0.84513223,  0.68757984],
       [ 0.90550232,  0.39838411,  0.88410946,  0.27235604,  0.83179862],
       [ 0.44378375,  0.37845376,  0.52914614,  0.38773563,  0.80787528],
       [ 0.44378375,  0.37845376,  0.52914614,  0.38773563,  0.80787528]])
...
>>> a[b,range(b.shape[0])]
array([ 0.4564633 ,  0.82383088,  0.88410946,  0.38773563,  0.80787528])

使用花式索引从numpy矩阵的每列中获取一个值

2 个答案: