python - 在 numpy 中创建下一个有效 bool 值的索引数组

Question

我有一个一维 numpy bool 数组，我的目标是获得一个相同维度的整数数组，其中每个条目都声明 bool 数组的下一个有效索引。这是一个例子:

import numpy as np
b=np.asarray((True, False, False, True, False),dtype=bool)
#It should be assigned such that in the end
#i=array([ 0,  3,  3,  3, -1])

是否有任何组合 numpy 函数可以在线性运行时且不使用循环的情况下实现此目的？

Answer 1

这是我能想到的一种方法:

import numpy as np
b = np.asarray((True, False, False, True, False), dtype=bool)
idx = np.r_[np.where(b)[0], -1]
res = idx[np.cumsum(np.r_[False, b[:-1]])]
print(res)
# [ 0  3  3  3 -1]

如果b非常大，您可以使用以下方式保存一个串联:

import numpy as np
b = np.asarray((True, False, False, True, False), dtype=bool)
idx = np.r_[np.where(b)[0], -1]
c = np.zeros(len(b), dtype=np.int32)
np.cumsum(b[:-1], out=c[1:])
res = idx[c]
print(res)
# [ 0  3  3  3 -1]