Sum rows of a 2D array with a specific stepsize - NumPy

时间:2017-12-22 10:09:57

标签: numpy

This is a quick one. I am wondering if there is a better way to express the following lines (besides using a short loop):

energy = np.zeros((4, signal.shape[1]))
energy[0::4, 0:] = np.sum(signal[0::4, :], axis=0)
energy[1::4, 0:] = np.sum(signal[1::4, :], axis=0)
energy[2::4, 0:] = np.sum(signal[2::4, :], axis=0)
energy[3::4, 0:] = np.sum(signal[3::4, :], axis=0)

1 个答案:

答案 0 :(得分:1)

重塑以将第一个轴分成两个,然后沿着这两个中的第一个相加,就像这样 -

energy = signal.reshape(-1,4,signal.shape[1]).sum(0)

示例运行 -

In [327]: np.random.seed(0)

In [328]: signal = np.random.randint(0,9,(8,5))

In [329]: energy = np.zeros((4, signal.shape[1]))
     ...: energy[0::4, 0:] = np.sum(signal[0::4, :], axis=0)
     ...: energy[1::4, 0:] = np.sum(signal[1::4, :], axis=0)
     ...: energy[2::4, 0:] = np.sum(signal[2::4, :], axis=0)
     ...: energy[3::4, 0:] = np.sum(signal[3::4, :], axis=0)

In [330]: energy
Out[330]: 
array([[ 13.,   4.,   6.,   3.,  10.],
       [  8.,   5.,   4.,   7.,  15.],
       [  7.,  11.,  11.,   4.,  13.],
       [  7.,   8.,   8.,   5.,  12.]])

In [331]: signal.reshape(-1,4,signal.shape[1]).sum(0)
Out[331]: 
array([[13,  4,  6,  3, 10],
       [ 8,  5,  4,  7, 15],
       [ 7, 11, 11,  4, 13],
       [ 7,  8,  8,  5, 12]])

对于行数不一定是4的倍数的数组,这里是通用版本 -

m = signal.shape[0]
n = m//4
energy = signal[:n*4].reshape(n,4,-1).sum(0)
energy[:m%4] += signal[n*4:]
相关问题