[FIXED] Vectorize the procedure of finding max values for splited numpy sub-arrays

Issue

Given a 1d numpy array arr and an indices array ind,

numpy.split(arr, ind)

returns the list of divided sub-arrays of arr by ind. We may use the Python for loop to get the max values for all the sub-arrays.

For example,

import numpy as np
arr = np.arange(12)
ind = np.array([3, 5, 9])
sub_arrays = np.split(arr, ind) 
# [array([0, 1, 2]), array([3, 4]), array([5, 6, 7, 8]), array([ 9, 10, 11])]
max_values = [sub_array.max() for sub_array in sub_arrays]
# [2, 4, 8, 11]

Can we do this in a numpy-style, vectorized way? For example, without splitting arr at first.

Solution

You can use np.maximum.reduceat. To get the exact result, you need to have ind also include index 0.

import numpy as np

arr = np.arange(12)
ind = np.array([3, 5, 9])
max_values = np.maximum.reduceat(arr, np.concatenate(([0], ind)))
print(max_values) # [ 2  4  8 11]

Answered By - jared