有没有for循环的替代方法吗?

时间:2020-08-02 16:25:34

标签: python arrays numpy

我有两个numpy数组AlA的尺寸为(n, x, y)l的尺寸为(n,1)。我得到的结果如下:

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

对于较大的n,这种获取结果的方法非常耗时。是否有其他方法可以快速获得相同的结果?

2 个答案:

答案 0 :(得分:2)

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

答案 1 :(得分:0)

列表理解是什么?

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