如何用NumPy向量化实现基于条件的for循环逻辑

背景 #

我有两个NumPy数组l1和l2，生成代码如下：

import numpy as np

start, jump = 1, 2
L = 8

l1 = np.arange(start, L)
l2 = np.arange(start + jump, L+jump)

生成的数组为：

l1 = [1 2 3 4 5 6 7]
l2 = [3 4 5 6 7 8 9]

需求 #

我需要生成两个结果数组r1和r2，规则是：逐个将l1和l2的元素分别添加到r1和r2时，需检查l1的第i个元素不在当前r2中，满足条件才执行添加。

我已经用Python循环实现了这个逻辑：

r1 = []
r2 = []

for i in range(len(l1)):
    if (l1[i] not in r2):
        r1.append(l1[i])
        r2.append(l2[i])

预期结果为：

r1 = [1, 2, 5, 6]
r2 = [3, 4, 7, 8]

作为NumPy新手，我想仅用NumPy实现该逻辑，不使用循环，请求帮助。

纯NumPy实现方案 #

针对当前场景的无循环最优解 #

你的场景中，l1是连续递增步长为1的数组，l2是l1加上固定偏移量jump，可以通过分析元素规律直接构建筛选掩码，完全不需要循环：

import numpy as np

start, jump = 1, 2
L = 8

l1 = np.arange(start, L)
l2 = np.arange(start + jump, L+jump)

# 构建掩码：索引i满足 i % (2*jump) < jump 时保留元素
mask = np.arange(len(l1)) % (2 * jump) < jump

# 用掩码筛选得到结果
r1 = l1[mask]
r2 = l2[mask]

print("r1 =", r1.tolist())
print("r2 =", r2.tolist())

输出：

r1 = [1, 2, 5, 6]
r2 = [3, 4, 7, 8]

规律说明 #

因为l2 = l1 + jump，当添加前jump个元素后，后续的jump个l1元素会正好等于已添加到r2的元素，因此需要跳过；再往后的jump个元素又不在当前r2中，可继续保留，以此循环。用索引取模的方式可以直接生成符合规则的掩码。

通用场景的NumPy优化方案（少量循环） #

如果你的数组不满足连续递增的规律，可以用NumPy的np.in1d函数结合少量循环实现，比原生Python循环效率更高：

import numpy as np

start, jump = 1, 2
L = 8

l1 = np.arange(start, L)
l2 = np.arange(start + jump, L+jump)

mask = np.zeros(len(l1), dtype=bool)
added_r2 = np.array([])

for i in range(len(l1)):
    # 用NumPy的in1d判断元素是否存在，比Python的in操作高效
    if not np.in1d(l1[i], added_r2):
        mask[i] = True
        added_r2 = np.append(added_r2, l2[i])

r1 = l1[mask]
r2 = l2[mask]

print("r1 =", r1.tolist())
print("r2 =", r2.tolist())

内容的提问来源于stack exchange，提问作者EducateMe