Question

假设我有一个NumPy数组A形状(N,N,N)。由此，我形成2D形状B (N,N)形状B = np.column_stack( tuple(A[i,:,i] for i in range(N)) )，如下所示：

i-th

换句话说，对于2D的{{1}} A切片，我会选择它i-th列;然后我stack这些列形成B。

我的问题是：

是否有更有效的方法（NumPy索引/切片）从B构建A;主要是，是否可以消除for 2D的{{1}}个切片上的内A循环？

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

对于所有i，有效地获得第i个2D NumPy阵列的2D切片列

1 个答案: