使用numpy生成由两个对象向量组成的矩阵对。

Question

我有两个对象数组，长度不一定相同：

import numpy as np

class Obj_A:
    def __init__(self,n):
        self.type = 'a'+str(n)
    def __eq__(self,other):
        return self.type==other.type

class Obj_B:
    def __init__(self,n):
        self.type = 'b'+str(n)
    def __eq__(self,other):
        return self.type==other.type

a = np.array([Obj_A(n) for n in range(2)])
b = np.array([Obj_B(n) for n in range(3)])

我想生成矩阵。

mat = np.array([[[a[0],b[0]],[a[0],b[1]],[a[0],b[2]]],
                [[a[1],b[0]],[a[1],b[1]],[a[1],b[2]]]])

这个矩阵的形状为(len(a),len(b),2)。它的元素是

mat[i,j] = [a[i],b[j]]

一个解决方案是：

mat = np.empty((len(a),len(b),2),dtype='object')
for i,aa in enumerate(a):
    for j,bb in enumerate(b): 
        mat[i,j] = np.array([aa,bb],dtype='object')

但是对于我的问题来说这太昂贵了，因为我的问题中 O(len(a)) = O(len(b)) = 1e5。

我怀疑有一个干净的numpy解决方案，涉及到np.repeat、np.tile和np.transpose，与这里的被接受的答案类似，但在这种情况下，输出不仅需要重塑以获得所需的结果。

- kevinkayaks

Answer 1

我建议使用np.meshgrid()，它接受两个输入数组，并沿不同轴重复这两个数组，这样查看输出的相应位置即可获得所有可能的组合。例如：

>>> x, y = np.meshgrid([1, 2, 3], [4, 5])
>>> x
array([[1, 2, 3],
       [1, 2, 3]])
>>> y
array([[4, 4, 4],
       [5, 5, 5]])

在您的情况下，您可以将这两个数组放在一起并将它们转置成适当的配置。根据一些实验，我认为这对您应该有效：

>>> np.transpose(np.meshgrid(a, b), (2, 1, 0))

- David Z

---

谢谢David。这个程序运行速度快，而且可以适应我感兴趣的问题规模。我不完全理解(2,1,0)是什么意思。有没有简单的方法来看看那里发生了什么？ - kevinkayaks

这只是告诉np.transpose()如何重新排列数组的维度。我建议查看文档并尝试一些示例，例如创建一个3D数组arr并尝试arr.transpose(0,2,1)，arr.transpose(2,0,1)等等，以了解其工作原理。 - David Z

对于一个3D数组，(2,1,0) 是默认的转置顺序，因此 np.transpose(np.meshgrid(a,b)) 就足够了。另一种表达方式是 np.array(np.meshgrid(a,b)).T。 - hpaulj