Increment and insert values based on a specific value of another array

Question

I have the following code with the given arrays a and b.

import numpy as np

# Parts of interest are highlighted with ^ ...

a = np.array([0,2,9,12,18,19])
#                   ^^    ^^
b = np.array([1,1,1,2,1,3]
#                   ^   ^
# Should result in an array like
assert result == np.array([0,2,9,12,13,18,19,20,21])
#                                ^^ ^^    ^^ ^^ ^^

The values in b define how many increments of the value in a (at the same index) should be inserted in the result. Ones in b don't affect the result. I think that I could do some splitting/joining and use a loop. But I'm wondering if this can be solved with numpy functions and good performance?

Thank you for your help!

Answer 1

Approach #1 : Here's a vectorized one -

def intervaled_ranges(ar, start=0):
    # Vectorized way to create ranges given sizes for each group
    c = ar.cumsum()
    v = -ar+1
    l = ar.sum()

    i = np.ones(l, dtype=int)
    i[c[:-1]] = v[:-1]
    i[0] = start
    return i.cumsum()

out = np.repeat(a,b)+intervaled_ranges(b)

Approach #2 : We can incorporate a into the intervaled-formation and hence skip the repeat step and achieve better performance, like so -

c = b.cumsum()
v = -b+1
s = b.sum()
i = np.ones(s, dtype=a.dtype)
i[c[:-1]] = v[:-1]+np.diff(a)
i[0] = a[0]
out = i.cumsum()

Answer 2

You can use a listcomp:

[k for i, j in zip(a, b) for k in range(i, i + j)]
# [0, 2, 9, 12, 13, 18, 19, 20, 21]

Answer 3

If pandas is an option:

d = pd.DataFrame({'a':a})
(d['a'].add(d.loc[d.index.repeat(b)]
            .groupby(level=0)
            .cumcount())
       .values
)

Output:

array([ 0,  2,  9, 12, 13, 18, 19, 20, 21], dtype=int64)