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