如何从numpy中的两个向量中获得所有对的笛卡尔积？

Question

这与this question 几乎相同，但在 numpy 中，以及针对某些维度的矩阵：

我有两个长度为 n 的矩阵，例如 A 和 B，

n = 1000
A = np.random.rand(n, 3, 4, 5)
B = np.random.rand(n, 3, 4, 5)

我想要一个新矩阵 C，其中每个元素 (C[i]) 都是形状为 (2, 3, 4, 5) 的矩阵：C[i][0] 是 A[k]，C[i][0] 是 B[k] 对于某些 @987654334 @，其中来自A 和B 的元素为另一个中的每个其他元素选择一次。

一个更具体的例子：

A = [elem1, elem2, elem3]
B = [elem4, elem5, elem6]

然后

C = [
[elem1, elem4],[elem1, elem5],[elem1, elem6],
[elem2, elem4],[elem2, elem5],[elem2, elem6],
[elem3, elem4],[elem3, elem5],[elem3, elem6]
]

我知道我可以使用 itertools 来构建它，但我认为它已经存在，而且速度更快。

Answer 1

从更简单的理解角度来看，基于作业的分配是有意义的。因此-

dt = np.result_type(A.dtype,B.dtype)
s = A.shape
C = np.empty((s[0],)+(2,)+s[1:],dtype=dt)
C[:,0] = A
C[:,1] = B

这可以通过单线的堆叠操作来简化 -

C = np.stack((A,B),axis=1)

解释一下这两种方法是如何解决的：

单线基本上是在 B 之后堆叠 A。在第一种方法中，结果是沿第二个轴的 2 元素 ndarray - s[0],)+(2,)+s[1:]。因此，离开 A 和离开 B 的每个元素都沿该轴交错。这个堆叠操作有一个 NumPy 内置的 np.stack，这是第二种方法。因此，这两个是等价的，可以解决我们的问题。

还有一些验证：

问题要求 - a new matrix C, in which every element (C[i]) is a matrix of shape (2, 3, 4, 5)。因此，C 的形状将是 (n,2,3,4,5)。堆叠沿第二轴。此外，据我了解，问题是要求C[k][0] == A[k] 为k=0..n 和C[k][1] == B[k] 为k=0..n。让我们证明-

In [323]: n = 1000
     ...: A = np.random.rand(n, 3, 4, 5)
     ...: B = np.random.rand(n, 3, 4, 5)

In [324]: C = np.stack((A,B),axis=1)

In [325]: np.all([np.allclose(C[i][0],A[i]) for i in range(1000)])
Out[325]: True

In [326]: np.all([np.allclose(C[i][1],B[i]) for i in range(1000)])
Out[326]: True