numpy 嵌套for循环的可能向量化

kupeojn6  于 2023-10-19  发布在  其他
关注(0)|答案(2)|浏览(77)

考虑两个相同的3D数组a1a2。在这个小例子中,我使用了形状为(10, 2, 2)的随机数组。我希望能够将vstack中的所有2x2 2D数组与a2中的所有2x2 2D数组进行a1,并生成一个包含堆叠的2D数组的3D数组。
因为我的实际a1a2阵列(总是具有相同的形状)的大小通常约为(100000, 20, 20),所以我通常选择只研究vstack艾德2D阵列的一个小子集。出于这个原因,我只研究了a1a2中的小的连续2D数组集,这些数组集由变量a1_starta1_stopa2_starta2_stop指定。
下面的代码工作正常,但我仍然需要做一些清理等:

import numpy as np

a1 = np.random.randint(0, 10, size=(10, 2, 2))
a2 = np.random.randint(0, 10, size=(10, 2, 2))

a1_start = 3
a1_stop = 5
a2_start = 0
a2_stop = 3

out = np.zeros(((a1_stop - a1_start)*(a2_stop - a2_start), 2*2, 2), dtype=int)
n=0

for p in a1[a1_start: a1_stop]:
    for q in a2[a2_start: a2_stop]:
        out[n] = np.vstack((p, q))
        n += 1
out

问题:因为我的实际数组相当大,* 我想知道上面的代码是否可以矢量化 *。我希望能够比较执行速度。

xytpbqjk

xytpbqjk1#

棘手:

arr = np.dstack(np.broadcast_arrays(
    a1[a1_start:a1_stop, None],
    a2[None, a2_start:a2_stop]
))
out = arr.reshape(-1, *arr.shape[2:])
c3frrgcw

c3frrgcw2#

你可以试试这样的东西:

q_range = np.arange(a2_start, a2_stop)
p_range = np.arange(a1_start, a1_stop)
p, q = np.meshgrid(p_range, q_range, indexing='ij')
q = q.flatten()
p = p.flatten()
out = np.concatenate((a1[p], a2[q]), axis=1)

相关问题