使用numpy从两个对象向量生成对矩阵

时间:2018-08-18 04:10:38

标签: python arrays numpy

我有两个对象数组,它们的长度不一定相同:

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.repeatnp.tilenp.transpose,类似于接受的答案here,但这种情况下的输出结果并不简单重塑为所需的结果。

1 个答案:

答案 0 :(得分:3)

我建议使用np.meshgrid(),它使用两个输入数组并沿不同的轴重复两个,以便查看输出的相应位置可以得到所有可能的组合。例如:

[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40], [51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90]]

在您的情况下,您可以将两个阵列放在一起并将它们转置为正确的配置。根据一些实验,我认为这应该对您有用:

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