Есть ли другой способ сделать то же самое без l oop?

Question

У меня есть два массива numpy A и l. Размер A равен (n, x, y), а размер l равен (n,1). Я получаю следующий результат:

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

Этот способ получения результата требует очень много времени для большего значения n. Есть ли альтернатива для быстрого получения такого же результата?

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])

Answer 2

как насчет понимания списка?

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