Python 如何从没有嵌套循环的二维阵列切片三维阵列？

我有四个矩阵，部分由彼此创建：

A是表示灰度图像堆栈的3D矩阵，形状为n、h、w B是也表示图像堆栈的3D矩阵，其中每个切片分别从a中的对应切片计算，并且也是形状n、h、w C是2D矩阵，包含z方向上最大值为B的索引，形状为h，w D是2D矩阵，其中a中的值从某个切片复制而来，该切片由位置x、y处C中的值表示。 使用循环实现的最小示例如下所示：

import numpy as np

A = np.random.randint(0, 255, size=(3, 4, 5))
B = np.random.randint(0, 255, size=(3, 4, 5))
C = np.argmax(B, axis=0)

D = np.zeros(C.shape, dtype=int)
for y in range(C.shape[0]):
    for x in range(C.shape[1]):
        D[y, x] = A[C[y, x], y, x]


> A
array([[[ 24,  84, 155,   8, 147],
        [ 25,   4,  49, 195,  57],
        [ 93,  76, 233,  83, 177],
        [ 70, 211, 201, 132, 239]],

       [[177, 144, 247, 251, 207],
        [242, 148,  28,  40, 105],
        [181,  28, 132,  94, 196],
        [166, 121,  72,  14,  14]],

       [[ 55, 254, 140, 142,  14],
        [112,  28,  85, 112, 145],
        [ 16,  72,  16, 248, 179],
        [160, 235, 225,  14, 211]]])

> B
array([[[246,  14,  55, 163, 161],
        [  3, 152, 128, 104, 203],
        [ 43, 145,  59, 169, 242],
        [106, 169,  31, 222, 240]],

       [[ 41,  26, 239,  25,  65],
        [ 47, 252, 205, 210, 138],
        [194,  64, 135, 127, 101],
        [ 63, 208, 179, 137,  59]],

       [[112, 156, 183,  23, 253],
        [ 35,   6, 233,  42, 100],
        [ 66, 119, 102, 217,  64],
        [ 82,  67, 135,   6,   8]]])

> C
array([[0, 2, 1, 0, 2],
       [1, 1, 2, 1, 0],
       [1, 0, 1, 2, 0],
       [0, 1, 1, 0, 0]])

> D
array([[ 24, 254, 247,   8,  14],
       [242, 148,  85,  40,  57],
       [181,  76, 132, 248, 177],
       [ 70, 121,  72, 132, 239]])

我的问题是：如何用C高效地分割A，消除嵌套的for循环？我最初的想法是将C扩展到一个3D布尔掩码，其中只有位置[C，y，x]被设置为真，然后简单地将其元素乘以a，并在z轴上求和。但是我想不出没有循环的pythonesque实现，如果我知道如何编写的话，我可能就不需要再创建布尔掩码了

---

我发现的最接近已经实现的功能是，但C只需要32个元素。

这里的标准方法是结合核心思想使用，也在以下内容中介绍：

y, x = np.ogrid[:C.shape[0],:C.shape[1]]
D = A[C[y, x], y, x]

将提出的方法与显式循环和基于np.ogrid的方法进行比较，可以得出：

import numpy as np


def take_by_axis_loop_fix(arr, pos):
    result = np.zeros(pos.shape, dtype=int)
    for i in range(pos.shape[0]):
        for j in range(pos.shape[1]):
            result[i, j] = arr[pos[i, j], i, j]
    return result


def take_by_axis_ogrid_fix(arr, pos):
    i, j = np.ogrid[:pos.shape[0], :pos.shape[1]]
    return arr[pos[i, j], i, j]


def take_by_axis_np(arr, pos, axis=0):
    return np.take_along_axis(arr, np.expand_dims(pos, axis=axis), axis=axis).squeeze()


def take_by_axis_ogrid(arr, pos, axis=0):
    ij = tuple(np.ogrid[tuple(slice(None, d, None) for d in pos.shape)])
    ij = ij[:axis] + (pos[ij],) + ij[axis:]
    return arr[ij]

表明这是迄今为止提出的最有效的方法。 还请注意，基于np.take_沿_轴的方法和基于_轴的take_ogrid方法对于高维输入基本不变，与_fix方法相反

---

特别是，take_by_axis_ogrid是take_by_axis_ogrid_fix的轴不可知版本，本质上是。

基本上，给定A、B和C，您正在寻找一种矢量化的无显式循环方法来计算D。这正确吗？是的，这正是我想要的！我已经看到了“沿轴移动”功能，但无法尝试。我的numpy版本似乎太旧了，由于不兼容的风险，我不敢升级。不过，谢谢你帮我试了一下，并给我定了时间！你有哪个版本的numpy？np.apply_沿_轴/np.take对您可用吗？我的numpy版本是1.14.5-我可以同时调用np.apply_沿_轴和np.take。这两个函数我都见过，但我不知道如何在这里使用它们。这正是我一直在寻找的，即使它可能不是最快的版本。非常感谢！不客气。根据我的经验，可读性比效率更重要。

import numpy as np


def take_by_axis_loop_fix(arr, pos):
    result = np.zeros(pos.shape, dtype=int)
    for i in range(pos.shape[0]):
        for j in range(pos.shape[1]):
            result[i, j] = arr[pos[i, j], i, j]
    return result


def take_by_axis_ogrid_fix(arr, pos):
    i, j = np.ogrid[:pos.shape[0], :pos.shape[1]]
    return arr[pos[i, j], i, j]


def take_by_axis_np(arr, pos, axis=0):
    return np.take_along_axis(arr, np.expand_dims(pos, axis=axis), axis=axis).squeeze()


def take_by_axis_ogrid(arr, pos, axis=0):
    ij = tuple(np.ogrid[tuple(slice(None, d, None) for d in pos.shape)])
    ij = ij[:axis] + (pos[ij],) + ij[axis:]
    return arr[ij]

A_ = np.random.randint(0, 255, size=(300, 400, 500))
B_ = np.random.randint(0, 255, size=(300, 400, 500))
C_ = np.argmax(B_, axis=0)

funcs = take_by_axis_loop_fix, take_by_axis_ogrid_fix, take_by_axis_ogrid, take_by_axis_np
for func in funcs:
    print(func.__name__, np.all(func(A_, C_) == take_by_axis_loop_fix(A_, C_)))
    %timeit func(A_, C_)
    print()

# take_by_axis_loop_fix True
# 10 loops, best of 3: 114 ms per loop

# take_by_axis_ogrid_fix True
# 100 loops, best of 3: 5.94 ms per loop

# take_by_axis_ogrid True
# 100 loops, best of 3: 5.54 ms per loop

# take_by_axis_np True
# 100 loops, best of 3: 3.34 ms per loop