[英]How to get the index of non-zero value in a numpy array quickly?
現在,我正在編寫一個函數,該函數將使用以下規則獲取非零值的索引:
[0,0,0,1,1,1,0,1,1,0]
的列表,它應該獲得列表[[3,4,5], [7,8]]
[0,0,1,1,1,2,2,1,1,0]
列表,預期結果為[[2,3,4],[5,6],[7,8]]
。 你有什么主意嗎? 先感謝您!
使用arr
作為輸入數組並具有數組列表作為輸出,您可以執行以下操作-
# Store non-zero element indices
idx = np.where(arr)[0]
# Get indices where the shifts occur, i.e. positions where groups of identical
# elements are separated. For this we perform differnetiation and look for
# non-zero values and then get those positions. Finally, add 1 to compensate
# for differentiation that would have decreased those shift indices by 1.
shift_idx = np.where(np.diff(arr[idx])!=0)[0]+1
# Split the non-zero indices at those shifts for final output
out = np.split(idx,shift_idx)
樣本輸入,輸出-
In [35]: arr
Out[35]: array([0, 0, 1, 1, 1, 2, 2, 1, 1, 0, 2, 2, 4, 3, 3, 3, 0])
In [36]: out
Out[36]:
[array([2, 3, 4]),
array([5, 6]),
array([7, 8]),
array([10, 11]),
array([12]),
array([13, 14, 15])]
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.