如何用NumPy表达式为满足条件的数组位置分配连续编号？

原代码的功能是创建与数组a同形状的数组b：在a中满足val < 0.666的位置填充从0开始的连续自然数，其余位置保留-1。完全可以用NumPy的向量化操作替代循环，实现更简洁高效的代码：

import numpy as np

a = np.random.rand(8, 8)
mask = a < 0.666
b = np.where(mask, np.cumsum(mask) - 1, -1)

代码解释： #

• mask = a < 0.666：生成一个布尔数组，标记a中所有满足条件的位置
• np.cumsum(mask)：对布尔数组做累积求和操作，满足条件的位置会依次得到1、2、3...的递增数值
• 减1操作np.cumsum(mask) - 1：让编号从0开始，和原代码的计数逻辑保持一致
• np.where(mask, ..., -1)：根据掩码数组选择对应值，满足条件的位置取累积和减1的结果，其余位置填充-1

如果需要替换其他条件，只需要修改mask的定义即可，比如复合条件mask = (a > 0.2) & (a < 0.7)也能直接适用。

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