python - 沿任意维度切片 numpy 数组

Question

假设我有一个 (40,20,30) numpy 数组，并且我有一个函数，在一些工作之后将沿着选定的输入轴返回输入数组的一半。有自动的方法吗？我想避免这样丑陋的代码:

def my_function(array,axis=0):

    ...

    if axis == 0:
        return array[:array.shape[0]/2,:,:] --> (20,20,30) array
    elif axis = 1:
        return array[:,:array.shape[1]/2,:] --> (40,10,30) array
    elif axis = 2: 
        return array[:,:,:array.shape[2]/2] --> (40,20,15) array

谢谢你的帮助

埃里克

Answer 1

我认为您可以为此使用 np.split [docs] , 并简单地获取返回的第一个或第二个元素，这取决于你想要哪个。例如:

>>> a = np.random.random((40,20,30))
>>> np.split(a, 2, axis=0)[0].shape
(20, 20, 30)
>>> np.split(a, 2, axis=1)[0].shape
(40, 10, 30)
>>> np.split(a, 2, axis=2)[0].shape
(40, 20, 15)
>>> (np.split(a, 2, axis=0)[0] == a[:a.shape[0]/2, :,:]).all()
True