python - 是否有另一种方法可以在没有 for 循环的情况下做同样的事情？

Question

我有两个 numpy 数组A和l. is 的维数和Ais(n, x, y)的维l数(n,1)。我得到的结果如下：

res = []
for i in range(n):
    res.append(A[i, x, l[i])

对于较大的值，这种获取结果的方式非常耗时n。是否有替代方法可以快速获得相同的结果？

标签： pythonarraysnumpy

Answer 1

如果0<=l[i]<y对于 的所有值i：

>>> n,x,y = 4,5,6
>>> A = np.random.randint(0,10,(n,x,y))
array([[[3, 3, 3, 8, 7, 0],
        [8, 1, 1, 5, 3, 8],
        [0, 1, 0, 4, 1, 3],
        [2, 2, 1, 8, 6, 5],
        [2, 5, 9, 2, 6, 3]],

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

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

       [[3, 8, 4, 3, 5, 6],
        [5, 3, 7, 7, 4, 6],
        [9, 0, 7, 9, 2, 1],
        [1, 6, 2, 2, 9, 5],
        [5, 0, 9, 0, 5, 2]]])
>>> l = np.random.randint(low=0, high=y-1, size=(n,1))
array([[0],
       [1],
       [3],
       [1]])
>>> x0 = 2
>>> res = []
>>> for i in range(n):
        res.append(A[i, x0, l[i])

>>> res
[array([0]), array([0]), array([4]), array([0])]

numpy：

>>> A[range(n), 2, l.flatten()]
array([0, 0, 4, 0])