与numpy rep一起使用的两个数组的元素编织

Question

我有两个不等长的数组，我正试图以一种特定的方式编织在一起，这是由等长数组mult1和mult2定义的。一般来说，我的数组很长（~1e6个元素），这是我计算中的一个性能关键瓶颈，因此我无法负担python for loop，因此我试图利用Numpy中的向量化函数。为了明确起见：

mult1 = np.array([0, 1, 2, 1, 0])
mult2 = np.array([1, 0, 1, 1, 0])

val1 = np.array([1, 2, 3, 4])
val2 = np.array([-1, -2, -3])

desired_final_result = np.array([-1, 1, 2, 3, -2, 4, -3])

val1和val2的编织是通过以下元素顺序通过mult1和mult2来定义的。两个mult数组的每个条目定义从相应的val数组中选择多少个元素。我们通过mult数组进行元素级处理；mult1[i]的值决定了我们从val1中选择多少个条目，然后我们继续使用mult2[i]的值来选择适当数量的val2条目的val1项，始终为每个索引i首先选择val1项

注意len（val1）=mult1.sum（）和len（val2）=mult2.sum（），因此我们总是得到一个带有len（期望的最终结果）=len（val1）+len（val2）的最终数组。在

极小例的显式解释

• 由于mult1的条目i=0是0，我们从val1中选择0条目的mult2的条目i=0，即1，因此我们从val2中选择1条目。这就解释了为什么期望的“最终”结果的第一个条目是-1。

• 由于mult1的条目i=1是1，我们从val1中选择1条目，并继续到mult2的条目i=1，即0，因此从val2中选择0个条目。这解释了为什么期望的“最终”结果的第二个条目是1。

• 由于mult1的条目i=2是2，我们从val1中选择下一个2条目的mult2的条目i=2，即1，因此我们从val2中选择下一个1条目。这解释了为什么期望的“最终”结果的条目2-4是2，3，-2。

• 由于mult1的条目i=3是1，所以我们从val1中选择下一个1条目，并继续到mult2的条目i=3，因此我们从val2中选择下一个1条目。这解释了为什么期望的最终结果的条目5-6是4，-3。

• 最后，由于mult1和mult2的i=4是0，所以我们没有什么可做的了，我们的数组被填满了。

问题

有没有办法使用矢量化函数，比如np.重复和/或np.选择解决我的问题？或者我需要用C编写这个计算并将其包装成python吗？在

Answer 1

这可以用裸体套路来完成，但我想到的最好方法是相当笨拙：

reps = numpy.empty([len(mult1)*2], dtype=int)
reps[::2] = mult1
reps[1::2] = mult2

to_repeat = numpy.empty_like(reps)
to_repeat[::2] = -1   # Avoid using 0 and 1 in case either of val1 or val2 is empty
to_repeat[1::2] = -2

indices = numpy.repeat(to_repeat, reps)
indices[indices==-1] = numpy.arange(len(val1))
indices[indices==-2] = numpy.arange(len(val1), len(val1) + len(val2))

final_result = numpy.concatenate([val1, val2])[indices]

Answer 2

在结果数组中创建布尔索引：

mult = np.array([mult1, mult2]).ravel('F')
tftf = np.tile([True, False], len(mult1))
mask = np.repeat(tftf, mult)

result = np.empty(len(val1) + len(val2), int)
result[ mask] = val1
result[~mask] = val2

编辑-我相信这也行得通：

^{pr2}$

它很短，但可能不会更快。在