工作中经常会用到Numpy操作,以下记录了几个小技巧,如掩码操作,保留非0元素,排序的索引值,等,这些能将代码写的很精简,尽量避免了for循环等低效操作。
a = np.array([0,0,6,3,4,0,0,5])
b = np.array([x!=0 for x in a]) # 掩码,0为FALSE,非0为TRUE
# [False False True True True False False True]
c = b + 0 # 将False/True转换为0/1
#[0 0 1 1 1 0 0 1]
print(a[b]) # 保留非0元素 [6 3 4 5]
print(a[b].argsort()) # 排序后的索引值 [1 2 3 0]