python - 按索引复制 numpy 数组的单个轴

Question

我正在寻找一种优雅的方法来通过索引提取 numpy 数组的单个轴的值。例如:

x = np.arange(16).reshape((4,4))
a = x[0]
b = x[:, 0] 

这是我通常做的事情，但是我正在寻找类似的东西:

a = get( x, axis=0, index=0)
b = get( x, axis=1, index=0)

是否有一些奇特的功能可以做到这一点？

Answer 1

您可以使用np.rollaxis将您感兴趣的轴移到前面，然后像平常一样对其进行索引:

def get(x, axis=0, index=0):
    return np.rollaxis(x, axis, 0)[index]

x = np.arange(27).reshape(3, 3, 3)

assert np.all(get(x, 1, 2) == x[:, 2, :])

正如 Joe 正确指出的那样，这将返回 x 的 View 。为了强制进行复制，您可以使用 .copy() 方法:

cpy = get(x, 1, 2).copy()