Skip to content
Advertisement

Consecutively split an array by the next max value

Suppose I have an array (the elements can be floats also):

D = np.array([0,0,600,160,0,1200,1800,0,1800,900,900,300,1400,1500,320,0,0,250])

The goal is, starting from the beginning of the array, to find the max value (the last one if there are several equal ones) and cut the anterior part of the array. Then consecutively repeat this procedure till the end of the array. So, the expected result would be:

[[0,0,600,160,0,1200,1800,0,1800],
 [900,900,300,1400,1500],
 [320],
 [0,0,250]]

I managed to find the last max value:

D_rev = D[::-1]
last_max_index = len(D_rev) - np.argmax(D_rev) - 1

i.e. I can get the first subarray of the desired answer. And then I can use a loop to get the rest.

My question is, if there is a numpy way to do it without looping?

Advertisement

Answer

IIUC, you can take the reverse cumulated max (see accumulate) of D to form groups, then split with itertools.groupby:

D = np.array([0,0,600,160,0,1200,1800,0,1800,900,900,300,1400,1500,320,0,0,250])

groups = np.maximum.accumulate(D[::-1])[::-1]
# array([1800, 1800, 1800, 1800, 1800, 1800, 1800, 1800, 1800, 1500, 1500,
#       1500, 1500, 1500,  320,  250,  250,  250])

from itertools import groupby
out = [list(list(zip(*g))[0]) for _, g in groupby(zip(D, groups), lambda x: x[1])]

# [[0, 0, 600, 160, 0, 1200, 1800, 0, 1800],
#  [900, 900, 300, 1400, 1500],
#  [320],
#  [0, 0, 250]]
User contributions licensed under: CC BY-SA
2 People found this is helpful
Advertisement