如何从任意形状numpy数组的最后一维提取值并解包赋值？

针对形状为任意s+(3,)的numpy数组A，要快速拆分出x=A[...,0]、y=A[...,1]、z=A[...,2]，可以用以下简洁方法：

通用最优解法 #

利用np.moveaxis将最后一维移至首位后直接解包：

import numpy as np

# 示例：任意形状的数组，比如(2, 4, 3)
A = np.random.rand(2, 4, 3)
x, y, z = np.moveaxis(A, -1, 0)

原理 #

np.moveaxis(A, -1, 0)会把数组A的最后一个维度（长度为3的维度）移动到第一个位置，将原形状s+(3,)转换为(3,)+s。此时直接解包，就能得到三个形状为s的数组，正好对应原数组的三个切片A[...,0]、A[...,1]、A[...,2]。

二维数组简化写法 #

如果数组是二维的（形状为(n,3)），可以直接用转置解包：

A = np.array([[1,2,3], [4,5,6]])
x, y, z = A.T

转置后数组形状变为(3,n)，解包后每个变量对应原数组的每一列，和A[:,0]、A[:,1]、A[:,2]完全一致。

验证示例 #

# 高维数组验证
A_high = np.random.rand(3, 5, 2, 3)
x, y, z = np.moveaxis(A_high, -1, 0)
print(x.shape)  # 输出(3,5,2)，与A_high[...,0].shape一致
print(np.array_equal(x, A_high[...,0]))  # 输出True，验证结果正确

内容的提问来源于stack exchange，提问作者Quercus Robur