Search code examples
pythonnumpyindexingmaxcumulative-sum

How to get accumulative maximum indices with numpy in Python?


I'm trying to get x and y coordinates of all cumulative maximums in a vector. I wrote a naive for-loop but I have a feeling there is some way to do it with numpy that's more elegant.

Does anyone know of any functions, masking-technique, or one-liners in numpy that can do this?

The details should be described by the plot below:

# Generate data
x = np.abs(np.random.RandomState(0).normal(size=(100)))
y = np.linspace(0,1,x.size)
z = x*y

# Get maximiums
idx_max = list()
val_max = list()

current = 0
for i,v in enumerate(z):
    if v > current:
        idx_max.append(i)
        val_max.append(v)
        current = v

# Plot them
with plt.style.context("seaborn-white"):
    fig, ax = plt.subplots(figsize=(10,3), ncols=2)
    ax[0].plot(z)
    ax[1].plot(z, alpha=0.618)
    ax[1].scatter(idx_max, val_max, color="black", edgecolor="maroon", linewidth=1.618)

enter image description here


Solution

  • We could make use of np.maximum.accumulate -

    idx_max = np.flatnonzero(np.isclose(np.maximum.accumulate(z),z))[1:]
    val_max = z[idx_max]
    

    The basic idea being that this accumulative max value when compared with the current element indicates which are the element positions responsible for "uplifting" the running max value. The responsible ones are the indices that we need as idx_max. Since, we are working with floating-pt numbers, we need to incorporate some tolerance with that np.isclose.

    That [1:] part is because the starting current value was 0 and so was z[0]. So, that v > current part won't append that starting index into output. To be precisely correct about it, it should had been -

    current = 0
    idx = np.flatnonzero(np.isclose(np.maximum.accumulate(z),z))
    idx = idx[z[idx] > current]
    

    But, given the starting values, the assumption made earlier keeps our code easier/compact/short.