find max value in islands defined by other vector

Question

I have a vector of values vals, a same-dimension vector of frequencies freqs, and a set of frequency values pins.

I need to find the max values of vals within the corresponding interval around each pin (from pin-1 to pin+1). However, the intervals merge if they overlap (e.g., [1,2] and [0.5,1.5] become [0.5,2]).

I have a code that (I think) works, but I feel is not optimal at all:

import numpy as np

freqs = np.linspace(0, 20, 50)
vals = np.random.randint(100, size=(len(freqs), 1)).flatten()
print(freqs)
print(vals)

pins = [2, 6, 10, 11, 15, 15.2]
# find one interval for every pin and then sum to find final ones
islands = np.zeros((len(freqs), 1)).flatten()
for pin in pins:
    island = np.zeros((len(freqs), 1)).flatten()
    island[(freqs >= pin-1) * (freqs <= pin+1)] = 1
    islands += island
islands = np.array([1 if x>0 else 0 for x in islands])
print(islands)

maxs = []
k = 0
idxs = []
for i,x in enumerate(islands):
    if (x > 0) and (k == 0):  # island begins
        k += 1
        idxs.append(i)
    elif (x > 0) and (k > 0):  # island continues
        pass
    elif (x == 0) and (k > 0):  # island finishes
        idxs.append(i)
        maxs.append(np.max(vals[idxs[0]:idxs[1]]))
        k = 0
        idxs = []
        continue
print(maxs)
    

Answer 1

I feel is not optimal at all

Good instincts! There is room for vectorisation here.

• Don't represent pins as a list, but rather as a numpy array.
• Don't write any loops.
• Stop creating two-dimensional arrays if you intend to flatten them immediately after.
• Calculate your islands by broadcasting comparisons for each pin into a matrix, and then reducing it to a vector over the logical_or function.
• Find the boundaries of your islands by applying a discrete differential.
• Find the maximum values in each island by masking your value vector with a group matrix, and then applying max over the second axis.

Suggested

import numpy as np

# This needs a constant seed for your test to be repeatable
rand = np.random.default_rng(seed=0)

freqs = np.linspace(start=(0,), stop=(20,), num=50)
vals = rand.integers(low=0, high=100, size=(len(freqs),))
pins = np.array((2, 6, 10, 11, 15, 15.2))

# Matrix, for each pin, of booleans for which values are in range
all_islands = (freqs >= pins - 1) & (freqs <= pins + 1)
# Reduce to a vector, one entry per value
islands = np.logical_or.reduce(all_islands, dtype=int, axis=1, keepdims=True)

# Discrete differential to find island boundaries (vector)
diffs = np.diff(islands.T.astype(int))
# Value indices needed to construct group masks (vector)
val_idx = np.arange(len(vals))[np.newaxis, :]
# Start and end indices for each island (vectors)
_, island_starts = np.nonzero(diffs == 1)
_, island_ends = np.nonzero(diffs == -1)

# For each island, of group masks (matrix)
groups = (val_idx > island_starts[:, np.newaxis]) & (val_idx < island_ends[:, np.newaxis])
# Reduce over the second axis getting maximum values in each group
maxes = np.max(groups * vals, axis=1)

assert np.all(maxes == (30, 97, 85, 54))