根据另一个数组的特定值递增和插入值 [英] Increment and insert values based on a specific value of another array

问题描述

我使用给定数组a和b的以下代码.

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])
#                                ^^ ^^    ^^ ^^ ^^

b中的值定义应在结果中插入a中的值的增量(在同一索引处). b中的那些不影响结果.我认为我可以进行一些拆分/合并并使用循环.但是我想知道这是否可以通过numpy函数和良好的性能来解决?

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?

谢谢您的帮助！

推荐答案

方法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)

方法2::我们可以将a合并到间隔形式中，从而跳过repeat步骤并获得更好的性能，就像这样-

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()

这篇关于根据另一个数组的特定值递增和插入值的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！