Numpy - 沿锯齿状数组最后一个轴的笛卡尔积的乘积

Question

给定一个元素数量不等的嵌套列表，我想找到计算沿最后一个轴的笛卡尔积的乘积的最快方法。换句话说，首先计算所有子列表之间的笛卡尔积，然后找到所有组合的乘积。最后，我想将这些值插入到与原始输入具有相同大小/维度的矩阵中。作为一个额外的复杂性，我想用额外的 0 填充形状 (1, ) 的轴。例如：

example1 = [[1, 2], [3, 4], [5], [6], [7]]

应该得到

[[[[[ 630.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]
  [[[ 840.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]]
 [[[[1260.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]
  [[[1680.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]]]

形状 (2, 2, 2, 2, 2)，尽管没有填充的话它会是 (2, 2, 1, 1, 1)。

我的初始功能是：

def convert_nest_to_product_tensor(nest):

    # find indices to collect elements from
    combinations = list(itertools.product(*[range(len(l)) for l in nest]))

    # collect elements and then calculate product for every Cartesian product
    products = np.array(
        [np.product([nest[i][idx] for i, idx in enumerate(comb)]) for comb in combinations]
        )

    # pad tensor for axes of shape 1
    tensor_shape = [len(l) for l in nest]
    tensor_shape = tuple([axis_shape+1 if axis_shape==1 else axis_shape for axis_shape in tensor_shape])
    tensor = np.zeros(tensor_shape)

    # insert values
    for i, idx in enumerate(combinations):
        tensor[idx] = products[i]

    return tensor

但是，这需要一段时间，特别是我找到笛卡尔积的乘积的部分。我尝试使用 np.meshgrid + np.stack 替换该组件：

products = np.stack(np.meshgrid(*nest), axis=-1).reshape(-1, len(nest))
products = np.prod(products, axis=-1)

虽然我更快地获得了正确的值，但它们的输出顺序不正确：

[[[[[ 630.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]
  [[[1260.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]]
 [[[[ 840.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]
  [[[1680.    0.]
    [   0.    0.]]
   [[   0.    0.]
    [   0.    0.]]]]]

任何反馈非常感谢如何（快速）完成这项工作！

Answer 1

获取笛卡尔元组和乘积的简单方法：

In [10]: alist = list(itertools.product(*example1))
In [11]: alist
Out[11]: [(1, 3, 5, 6, 7), (1, 4, 5, 6, 7), (2, 3, 5, 6, 7), (2, 4, 5, 6, 7)]
In [12]: [np.prod(x) for x in alist]
Out[12]: [630, 840, 1260, 1680]

或者使用 math.prod 来获得非 numpy 解决方案。