假设我有9个大小为N的数组(A,B,C,... J)。我想创建一个新的N 3x3矩阵数组,例如。
matrices[i] = [[A[i], B[i], C[i]],
[D[i], E[i], F[i]],
[G[i], H[i], J[i]]]
一个简单的解决方案是在for循环中将每个条目添加到数组matrices
中:
for i in range(len(matrices)):
matrices[i] = [[A[i], B[i], C[i]],
[D[i], E[i], F[i]],
[G[i], H[i], J[i]]]
任何人都有一些关于如何以更快的矢量化方式完成此操作以避免for循环的技巧?如果存在一些智能索引操作或其他什么。
答案 0 :(得分:4)
一种方法是将这些列与np.column_stack
堆叠在一起,然后重新整形np.reshape
-
np.column_stack((A,B,C,D,E,F,G,H,J)).reshape(-1,3,3)
与np.concatenate
的连接速度要快得多,因此将其与2D transpose
一起使用并重新整形 -
np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(9,-1).T.reshape(-1,3,3)
另一个np.concatenate
,3D transpose
和重塑 -
np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(3,3,-1).transpose(2,0,1)
运行时测试 -
In [59]: # Setup input arrays
...: N = 1000
...: A = np.random.randint(0,9,(N,))
...: B = np.random.randint(0,9,(N,))
...: C = np.random.randint(0,9,(N,))
...: D = np.random.randint(0,9,(N,))
...: E = np.random.randint(0,9,(N,))
...: F = np.random.randint(0,9,(N,))
...: G = np.random.randint(0,9,(N,))
...: H = np.random.randint(0,9,(N,))
...: J = np.random.randint(0,9,(N,))
...:
In [60]: %timeit np.column_stack((A,B,C,D,E,F,G,H,J)).reshape(-1,3,3)
10000 loops, best of 3: 84.4 µs per loop
In [61]: %timeit np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(9,-1).T.reshape(-1,3,3)
100000 loops, best of 3: 15.8 µs per loop
In [62]: %timeit np.concatenate((A,B,C,D,E,F,G,H,J)).reshape(3,3,-1).transpose(2,0,1)
100000 loops, best of 3: 14.8 µs per loop