查找（并记录）numpy数组的切片的最大值

Question

Given a numpy array such as a = [0, 3, 2, 4, 0, 2, 3, 1, 1, 6, 2] , is there a simple way to record the maximum value for every 3 values? The length of the array may not be a multiple of 3. In this case, the result should be b = [3, 4, 3, 6] .

I thought of something along the lines of

b = [max(a[k:k+3]) for k in range(0, len(a), 3)

but it doesn't take into account the values after the last multiple of 3 (which it should).

I've also thought of rearranging the numpy array such that it has 3*n rows, and taking the maxima along the proper axis using the numpy module, but, again, I'm not sure how to deal with the values after the last multiple of 3.

Answer 1

Approach #1

We can use np.ufunc.reduceat for performing such grouped/intervaled reduction operations. Thus, to get maximum values within each interval, we would have -

W = 3 # group width
np.maximum.reduceat(a,np.r_[:len(a):W])

Sample run -

In [166]: a
Out[166]: array([0, 3, 2, 4, 0, 2, 3, 1, 1, 6, 2])

In [167]: W = 3

In [168]: np.maximum.reduceat(a,np.r_[:len(a):W])
Out[168]: array([3, 4, 3, 6])

Approach #2

Here's another with slicing -

def max_interval_slice(a, W=3):
    n = len(a)//W
    max0 = a[:n*W].reshape(-1,W).max(1)
    if n*W==len(a):
        return max0
    else:
        return np.r_[max0, np.max(a[n*W:])]

Sample runs -

# Input array of length NOT multiple of width=3
In [99]: a
Out[99]: array([0, 3, 2, 4, 0, 2, 3, 1, 1, 6, 2])

In [100]: max_interval_slice(a, W=3)
Out[100]: array([3, 4, 3, 6])

# Input array of length multiple of width=3
In [95]: a = a[:9]

In [96]: max_interval_slice(a, W=3)
Out[96]: array([3, 4, 3])

Answer 2

First use np.pad

a = np.pad(a, [0, 1], mode='constant')

Then reshape and max

>>> np.max(a.reshape(-1, 3), axis=1)
array([3, 4, 3, 6])

To generalize this, just calculate padding in order to reshape to the desired dimension.

Answer 3

To minimize the amount of reallocation you need to do, you can compute the maxima of all the elements that fit into a multiple of 3, and then the max of the remainder. This solution is not as straightforward, but it does not create unnecessary copies of the data:

n = 3  # The group width
prefix, suffix = divmod(a.size, n)
output = np.empty(prefix + bool(suffix))
a[:n * prefix].reshape(-1, n).max(axis=1, out=output[:prefix])
if suffix:
    output[-1] = a[-suffix:].max()