一种遍历一次即可得到TOPK的算法
def get_topk(tensor_1d, topk=3):
# value in topk_vals are placed by descending order
topk_vals = [-float("Inf")] * topk
topk_idxs = [0] * topk
for idx, elem in enumerate(tensor_1d):
if elem > topk_vals[topk - 1]:
for i in range(topk):
# find where current top value should be placed
# then we right shift the topk_vals to place the top value
if elem > topk_vals[i]:
# right shift
for j in reversed(range(i, topk-1)):
topk_vals[j+1] = topk_vals[j]
topk_idxs[j+1] = topk_idxs[j]
topk_vals[i] = elem
topk_idxs[i] = idx
break
return topk_vals, topk_idxs
tensor_1d=[1,2,3,4,4,5,5,6]
topk_vals, topk_idxs = get_topk(tensor_1d, topk=3)
print(topk_vals)
print(topk_idxs)
该方法保存top k元素的数组,然后遍历每个元素,依次向后与该数组元素比较,找到大于当前top元素的位置,然后从当前位置右移top k元素的数组并插入该top元素。
一种基于直方图的方法
有时候需要N个数选出M个最小或最大值算法,但并未要求选出的M个数据需要有序排列,那么这使得算法复杂度可以很低。这里给出一种可行的方法,比常见的一些方法具有更加显著的速度。
1,找出N个数据的最大和最小值。需要一次N个数据遍历。
2,根据最大最小值根据一个间隔创建一个直方图,N个数据遍历一次,进行直方图统计。例如,直方图间隔为k,则每次直方图(CurDat-MinDat)/K位置加1即可.
3,根据直方图从最小或最大处开始,找到累积大于等于M个元素的阈值,需要一次直方图遍历。
4,根据步骤3找到的阈值选出M1个数据,需要一次N个数据遍历。
5,由于M1可能略大于M,稍作后处理移除即可。可能需要对M1进行多余数据次数的排序。
该方法的缺点则是第3步缺乏准确性,导致选出的数据会略大于需求个数,在一些极端情况下可能工作不好。但在我的应用中,未限定死选出的最小数必须为M个,只需要近似即可,因此该方法具有较好实现性和性能。