python - 从多维数组中选择一些维度

Question

我有一个 4 维数组，我需要沿给定维度处理该数组中的所有 1 维向量。这很好用：

    def myfun(arr4d,selected_dim): # selected_dim 只能是 2 或 3
        打印 arr4d.shape # (2, 5, 10, 10)
        对于 xrange(arr4d.shape[0]) 中的 i：
            对于 xrange(arr4d.shape[1]) 中的 j：
                对于 xrange 中的 k（arr4d.shape[selected_dim]）：
                    如果 selected_dim==2：
                        arr=arr4d[i,j,k,:]
                    elif selected_dim==3：
                        arr=arr4d[i,j,:,k]
                    do_something(arr) # arr 是一维的，有 10 个项目

...但我相信有一些方法可以避免嵌套的“if”部分，也许也更有效？就像在循环之前创建这个数组的其他视图然后遍历这些视图？

Answer 1

处理此问题的一种常见方法是使用np.rollaxis：

def myfun(arr4d, selected_dim):  # selected_dim can only be 2 or 3
    arr4d = np.rollaxis(arr4d, selected_dim)
    print arr4d.shape   # (10, 2, 5, 10)
    for i in xrange(arr4d.shape[1]):
        for j in xrange(arr4d.shape[2]):
            for k in xrange(arr4d.shape[0]):
                arr=arr4d[k, i, j, :]
                do_something(arr)  # arr is 1-D and has 10 items

请注意，它np.rollaxis应该返回一个视图，因此它实际上并不复制数组。