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

【问题标题】：Increment and insert values based on a specific value of another array根据另一个数组的特定值递增和插入值
【发布时间】：2019-11-12 16:15:41
【问题描述】：

我有以下代码以及给定数组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 函数和良好的性能来解决？

感谢您的帮助！

【问题讨论】：

a 的值是否不同？例如，有没有可能a = [1,1], b=[2,2]？
@QuangHoang a 的值是唯一的（按升序排列），b 的值可能相同。

标签： python numpy vectorization

【解决方案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 合并到intervaled-formation 中，从而跳过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()

【讨论】：

【解决方案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]

【讨论】：

【解决方案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)

【讨论】：