使用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 ，类似于此处接受的答案，但在这种情况下的输出不会简单地重塑为所需的结果。

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))