python - 与 numpy repeat 一起使用的两个数组的逐元素编织

Question

我有两个长度不等的数组 val1 和 val2，我试图以等长数组 定义的特定方式将它们编织在一起mult1 和 mult2。一般来说，我的数组很长(~1e6 个元素)，这是我计算中的性能关键瓶颈，所以我无法承受 python-for 循环，所以我试图利用 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< 的索引的逐元素处理定义/em>。两个 mult 数组的每个条目都定义了从相应的 val 数组中选择多少个元素。我们通过 mult 数组逐元素进行； mult1[i] 的值决定了我们从 val1 中选择了多少条目，然后我们继续根据 mult2[i] 的值进行选择适当数量的 val2 条目，始终为每个索引 i 选择 val1 条目。

请注意 len(val1) = mult1.sum() 和 len(val2) = mult2.sum()，因此我们总是以最终数组结束len(desired_final_result) = len(val1) + len(val2)。

最小示例的明确解释

• 由于 mult1 的条目 i=0 是 0，我们从 < em>val1 并移动到 mult2 的条目 i=0，即 1，因此我们选择 1 来自 val2 的条目。这就解释了为什么 desired_final_result 的第一个条目是 -1。

• 由于 mult1 的条目 i=1 是 1，我们从 < em>val1 并移至 mult2 的条目 i=1，即 0，因此选择 0 来自 val2 的条目。这就解释了为什么 desired_final_result 的第二个条目是 1。

• 由于 mult1 的条目 i=2 是 2，我们选择接下来的 2 条目从 val1 移动到 mult2 的条目 i=2，即 1，因此我们选择下一个1 个条目来自 val2。这解释了为什么 desired_final_result 的条目 2-4 是 2、3、-2。

• 由于 mult1 的条目 i=3 是 1，我们选择下一个 1 条目从 val1 移动到 mult2 的条目 i=3，它也是 1，因此我们选择val2 的下一个 1 条目。这解释了为什么 desired_final_result 的条目 5-6 是 4, -3。

• 最后，由于 mult1 和 mult2 的 i=4 都是 0，我们无事可做，我们的数组已满。

问题

有没有办法使用矢量化函数(例如 np.repeat 和/或 np.choose)来解决我的问题？或者我是否需要求助于用 C 编写此计算并将其包装到 python 中？

Answer 1

在结果数组中创建 bool 索引:

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

编辑 - 我相信这也有效:

idx = np.repeat(mult1.cumsum(), mult2)
result = np.insert(val1, idx, val2)

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