如何用NumPy向量化优化0-1数组交替1的校验算法？

我有两个仅由0和1组成的数组：

a = [1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
b = [0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1]

要求从左到右遍历时，1必须在两个数组间交替出现，不能连续在同一个数组中出现两次。

我目前用Pandas结合循环实现了这个逻辑，但效率很低：

import pandas as pd

df = pd.DataFrame({"A": a, "B": b})
df2 = df[(df.A > 0) | (df.B > 0)]
i = 0
for idx in df2.index:
    try:       
        if df2.at[idx, 'A'] == df2.at[df2.index[i + 1], 'A']:
            df.at[idx, 'A'] = 0
        if df2.at[idx, 'B'] == df2.at[df2.index[i + 1], 'B']:
            df.at[idx, 'B'] = 0
        i += 1
    except IndexError:
        pass

请问如何用NumPy进行向量化处理来提升速度？

核心思路是先提取所有出现1的位置及对应的数组标识，筛选出符合交替规则的有效位置，最后将原数组中不符合规则的1置为0，全程用向量化操作替代循环，大幅提升效率。

具体步骤： #

1. 将原始列表转为NumPy数组，开启向量化操作基础：

import numpy as np

a_np = np.array(a)
b_np = np.array(b)

2. 提取所有1的位置并标记来源数组：
   用1标记来自数组a的1，2标记来自数组b的1，然后按原始顺序合并排序：

# 收集a中1的索引和对应标记
a_ones = np.where(a_np == 1)[0]
a_markers = np.full_like(a_ones, 1)
# 收集b中1的索引和对应标记
b_ones = np.where(b_np == 1)[0]
b_markers = np.full_like(b_ones, 2)

# 合并后按索引从小到大排序，保证顺序和原数组遍历一致
all_indices = np.concatenate([a_ones, b_ones])
all_markers = np.concatenate([a_markers, b_markers])
sort_idx = np.argsort(all_indices)
sorted_indices = all_indices[sort_idx]
sorted_markers = all_markers[sort_idx]

3. 生成有效掩码筛选符合规则的项：
   第一个1默认保留，从第二个开始，仅保留与前一个来源数组不同的1：

valid_mask = np.ones_like(sorted_markers, dtype=bool)
# 从第二个元素开始，判断是否与前一个标记不同
valid_mask[1:] = sorted_markers[1:] != sorted_markers[:-1]

# 提取有效位置和对应标记
valid_indices = sorted_indices[valid_mask]
valid_markers = sorted_markers[valid_mask]

4. 重置原数组，仅保留有效位置的1：

# 先清空原数组
a_np[:] = 0
b_np[:] = 0

# 给有效位置赋值1
a_np[valid_indices[valid_markers == 1]] = 1
b_np[valid_indices[valid_markers == 2]] = 1

完成后，a_np和b_np即为符合要求的数组。这种方案完全规避了循环，在数组规模较大时，性能会远优于Pandas循环实现。

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