python - 从 numpy 数组的最后一个轴提取值

Question

考虑一个形状为 (3,) 的 numpy 数组 A。然后是下面一行

x,y,z = A

将x分配给A[0]，将y分配给A[1]和z 到 A[2]。现在假设 A 的形状为 s + (3,) ，以表示任意形状 s。我想类似地将 x 分配给 A[...,0]，将 y 分配给 A[...,1 ] 和 z 到 A[...,2]。上面一行

x,y,z = A

不起作用并给出 ValueError: 没有足够的值来解包(预期为 3，得到 2)[当 A 具有形状 (2,3) 时]。我怎样才能以干净的方式完成所需的任务？显然如下

x,y,z = A[...,0], A[...,1], A[...,2]

可以工作，但如果将 3 替换为某个大数字，则有点乏味。

Answer 1

您可以使用numpy.rollaxis :

x, y, z = np.rollaxis(A, -1)

假设这个输入:

array([[[0, 1, 2],
        [3, 4, 5]]])

输出:

# x, y, z
(array([[0, 3]]), array([[1, 4]]), array([[2, 5]]))

这适用于任何位置，只需指定用作rollaxis第二个参数的尺寸即可:

x, y, z = np.rollaxis(A, 2)

时间:

在形状 = (7,6,8,2,1,4,5,9,3) ; A = np.arange(np.prod(shape)).reshape(shape) 作为输入。

%timeit x, y, z = np.rollaxis(A, -1)
# 1.96 µs ± 59 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

%timeit x, y, z = np.moveaxis(A, -1, 0)
# 3.78 µs ± 298 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

# credit to @QuangHoang
%timeit x, y, z = [A[...,i] for i in range(3)]
# 691 ns ± 8.24 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)