沿任意维度切割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