根据另一个数组的特定值递增和插入值

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 中的那些不影响结果。我认为我可以做一些 splitting/joining 并使用循环。但我想知道这是否可以通过 numpy 函数和良好的性能来解决？

感谢您的帮助！

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

如果 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)

Answer 3

您可以使用列表组件：

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

根据另一个数组的特定值递增和插入值

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

python

numpy

vectorization