Numpy如何对满足条件的子数组填充递增序列替代循环写法

你需要的功能可以用以下两种更简洁的原生Numpy方式实现，无需显式声明额外的中间变量：

方案1：纯一行无中间变量实现 #

适合判断逻辑计算成本低的场景，代码最紧凑：

b = np.where(predicate(a), np.cumsum(predicate(a)) - 1, -1)

原理说明：np.cumsum对布尔数组求和时会将True识别为1、False识别为0，累加后将结果减1刚好得到从0开始的递增值，np.where会在满足条件的位置填充累加值，其余位置填充-1。

方案2：Python 3.8+ 高效一行实现 #

如果不想重复执行两次判断逻辑，可以用海象运算符隐式存储中间布尔数组，效率和你原来的写法完全一致：

b = np.where((m := predicate(a)), np.arange(np.count_nonzero(m)), -1)

这个写法仅将你原来的三行代码合并为一行，没有额外性能损耗。

补充说明 #

如果你使用的Python版本低于3.8，那么你原来的写法已经是效率最优的实现：

• 中间布尔数组m的内存开销极低，布尔类型每个元素仅占1字节，即使千万级长度的数组也仅占用几MB内存
• 两次引用m的操作耗时可以忽略，对整体性能没有影响

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