沿任意轴的未知维度的numpy数组的连接

Question

我有两个数组A和B未知的维度,我想沿着这个N维度连接.例如:

>>> A = rand(2,2)       # just for illustration, dimensions should be unknown
>>> B = rand(2,2)       # idem
>>> N = 5

>>> C = concatenate((A, B), axis=N)
numpy.core._internal.AxisError: axis 5 is out of bounds for array of dimension 2

>>> C = stack((A, B), axis=N)
numpy.core._internal.AxisError: axis 5 is out of bounds for array of dimension 3

这里询问相关问题.不幸的是,当尺寸未知时,所提出的解决方案不起作用,我们可能必须添加几个新轴,直到获得最小尺寸N.

我所做的是将形状从1向上扩展到第Nth维,然后连接:

newshapeA = A.shape + (1,) * (N + 1 - A.ndim)
newshapeB = B.shape + (1,) * (N + 1 - B.ndim)
concatenate((A.reshape(newshapeA), B.reshape(newshapeB)), axis=N)

使用此代码,我应该能够将(2,2,1,3)数组与沿轴3的(2,2)数组连接起来.

有没有更好的方法来实现这一目标？

ps:根据建议更新第一个答案.

Answer 1

我不认为你的方法有什么问题，尽管你可以使你的代码更紧凑一些：

newshapeA = A.shape + (1,) * (N + 1 - A.ndim)