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中的那些不会影响结果。我认为我可以进行一些拆分/合并并使用循环。但是我想知道这是否可以通过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

您可以使用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)

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

3 个答案: