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

Question

I have two numpy arrays A and l . The dimension of A is (n, x, y) and the dimension of l is (n,1) . I get the result as follows:

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

This way of getting the result is very time consuming for a larger value of n . Is there an alternative to get the same result quickly?

Answer 1

If 0<=l[i]<y for all values of 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])

Answer 2

what about list comprehension?

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