group argmax/argmin over partitioning indices in numpy

Question

Numpy's ufuncs have a reduceat method which runs them over contiguous partitions within an array. So instead of writing:

import numpy as np
a = np.array([4, 0, 6, 8, 0, 9, 8, 5, 4, 9])
split_at = [4, 5]
maxima = [max(subarray for subarray in np.split(a, split_at)]

I can write:

maxima = np.maximum.reduceat(a, np.hstack([0, split_at]))

Both will return the maximum values in slices a[0:4], a[4:5], a[5:10], being [8, 0, 9].

I would like a similar function to perform argmax, noting that I would only like a single maximum index in each partition: [3, 4, 5] with the above a and split_at (despite indices 5 and 9 both obtaining the maximum value in the last group), as would be returned by

np.hstack([0, split_at]) + [np.argmax(subarray) for subarray in np.split(a, split_at)]

I will post a possible solution below, but would like to see one that is vectorized without creating an index over groups.

Answer 1

This solution involves building an index over groups ([0, 0, 0, 0, 1, 2, 2, 2, 2, 2] in the above example).

group_lengths = np.diff(np.hstack([0, split_at, len(a)]))
n_groups = len(group_lengths)
index = np.repeat(np.arange(n_groups), group_lengths)

Then we can use:

maxima = np.maximum.reduceat(a, np.hstack([0, split_at]))
all_argmax = np.flatnonzero(np.repeat(maxima, group_lengths) == a)
result = np.empty(len(group_lengths), dtype='i')
result[index[all_argmax[::-1]]] = all_argmax[::-1]

To get [3, 4, 5] in result. The [::-1]s ensure that we get the first rather than the last argmax in each group.

This relies on the fact that the last index in fancy assignment determines the value assigned, which @seberg says one shouldn't rely on (and a safer alternative can be achieved with result = all_argmax[np.unique(index[all_argmax], return_index=True)[1]], which involves a sort over len(maxima) ~ n_groups elements).