如何在NumPy中用不同形状数组扩展多维数组的最后维度

我有以下两个NumPy数组：

import numpy as np

A = np.array([
            [[[0, 1, 2, 3],
              [3, 0, 1, 2],
              [2, 3, 0, 1],
              [1, 3, 2, 1],
              [1, 2, 3, 0]]],
            
            [[[9, 8, 7, 6],
              [5, 4, 3, 2],
              [0, 9, 8, 3],
              [1, 9, 2, 3],
              [1, 0, -1, 2]]],
            
            [[[0, 7, 1, 2],
              [1, 2, 1, 0],
              [0, 2, 0, 7],
              [-1, 3, 0, 1],
              [1, 0, 1, 0]]]
              ])

print(A.shape)
# 输出: (3, 1, 5, 4)

B = np.array([
    [[[1, 0],
      [-1, 2],
      [9, 1],
      [8, 2],
      [7, 0]]],
            
      [[[9, 6],
       [5, 2],
       [0, 3],
       [1, 9],
       [1, 0]]],
            
      [[[0, 7],
        [1, 0],
        [0, 7],
        [-1, 1],
        [0, 0]]]
])

print(B.shape)
# 输出: (3, 1, 5, 2)

我需要将数组B沿数组A的最后一个维度进行拼接，得到如下结果数组X：

X = np.array([
            [[[0, 1, 2, 3, 1, 0],
              [3, 0, 1, 2,-1, 2],
              [2, 3, 0, 1, 9, 1],
              [1, 3, 2, 1, 8, 2],
              [1, 2, 3, 0, 7, 0]]],
            
            [[[9, 8, 7, 6, 9, 6],
              [5, 4, 3, 2, 5, 2],
              [0, 9, 8, 3, 0, 3],
              [1, 9, 2, 3, 1, 9],
              [1, 0,-1, 2, 1, 0]]],
            
            [[[0, 7, 1, 2, 0, 7],
              [1, 2, 1, 0, 1, 0],
              [0, 2, 0, 7, 0, 7],
              [-1,3, 0, 1,-1, 1],
              [1, 0, 1, 0, 0, 0]]]
              ])

print(X.shape)
# 输出: (3, 1, 5, 6)

实现方法 #

使用NumPy的np.concatenate()函数，指定axis=-1（代表最后一个维度）即可完成拼接：

import numpy as np

# 拼接数组A和B
X = np.concatenate([A, B], axis=-1)

# 验证结果维度
print(X.shape)  # 输出: (3, 1, 5, 6)

该函数会在指定维度上合并数组，由于A和B的前三个维度完全匹配（3,1,5），仅最后一个维度不同，因此可以直接沿最后维度拼接得到目标数组。

内容的提问来源于stack exchange，提问作者Amina Umar