從布爾數組中設置Python numpy索引

Question

如何將布爾數組轉換爲可迭代的索引？

例如，

import numpy as np 
import itertools as it 
x = np.array([1,0,1,1,0,0]) 
y = x > 0 
retval = [i for i, y_i in enumerate(y) if y_i] 

是否有更好的辦法嗎？

Answer 1

嘗試np.where或np.nonzero。

x = np.array([1, 0, 1, 1, 0, 0]) 
np.where(x)[0] # returns a tuple hence the [0], see help(np.where) 
# array([0, 2, 3]) 
x.nonzero()[0] # in this case, the same as above. 

見help(np.where)和help(np.nonzero)。

可能值得注意的是，在np.where頁面中提到，對於1D x而言，它基本上等同於問題中的longform。