Увеличение и вставка значений на основе определенного значения другого массива

Question

У меня есть следующий код с данными массивами a и 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])
#                                ^^ ^^    ^^ ^^ ^^

Значения в b определяют, сколько приращений значения в a (с тем же индексом) следует вставить в результат. Единицы в b не влияют на результат. Я думаю, что я мог бы сделать некоторые расщепление / соединение и использовать цикл. Но мне интересно, можно ли это решить с помощью бесчисленных функций и хорошей производительности?

Спасибо за помощь!

Answer 1

Подход № 1: Вот векторизованный -

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)

Подход № 2: Мы можем включить a в интервальную формацию и, следовательно,пропустите шаг repeat и добейтесь лучшей производительности, вот так -

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

Вы можете использовать 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

Если pandas является опцией:

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

Выход:

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