Home > Enterprise >  Consecutively split an array by the next max value
Consecutively split an array by the next max value

Time:03-03

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?

CodePudding user response:

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]]
  • Related