高效地获取 3D NumPy 数组的第 i 个二维切片的第 i 个列，对于所有 i

Question

假设我有一个 NumPy 数组 A，形状为 (N,N,N)。由此，我形成了一个形状为 (N,N) 的 2D 数组 B，如下所示：

B = np.column_stack( tuple(A[i,:,i] for i in range(N)) )

换句话说，对于 A 的 i-th 2D 部分，我认为它是 i-th 列；然后我 stack 这些列形成 B.

我的问题是：

有没有更有效的方法（NumPy indexing/slicing）从A构造B；主要是，是否可以消除 A 的 2D 切片上的内部 for 循环？

Answer 1

idx = np.arange(N)  # or idx = range(N)
A[idx,:,idx].T

示例：

import numpy as np
A = np.arange(27).reshape(3,3,3)

idx = np.arange(3)
A[idx,:,idx].T
#array([[ 0, 10, 20],
#       [ 3, 13, 23],
#       [ 6, 16, 26]])

np.column_stack( tuple(A[i,:,i] for i in range(3)) )
#array([[ 0, 10, 20],
#       [ 3, 13, 23],
#       [ 6, 16, 26]])

时序：大数组更快

def adv_index(N):
    idx = range(N)
    return A[idx,:,idx].T

N = 100
import numpy as np
A = np.arange(N*N*N).reshape(N,N,N)
    
%timeit np.column_stack(tuple(A[i,:,i] for i in range(N)))
# The slowest run took 4.01 times longer than the fastest. This could mean that an intermediate result is being cached.
# 1000 loops, best of 3: 210 µs per loop

%timeit adv_index(N)
# The slowest run took 5.87 times longer than the fastest. This could mean that an intermediate result is being cached.
# 10000 loops, best of 3: 51.1 µs per loop

(np.column_stack(tuple(A[i,:,i] for i in range(N))) == adv_index(N)).all()
# True

高效地获取 3D NumPy 数组的第 i 个二维切片的第 i 个列，对于所有 i

Efficiently getting i-th column of i-th 2D slice of 3D NumPy array, for all i

python

numpy

matrix-indexing