github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/exec/topn.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package exec
     6  
     7  import (
     8  	"container/heap"
     9  
    10  	"github.com/grailbio/base/log"
    11  )
    12  
    13  type topIndex struct{ Index, Count int }
    14  
    15  type topHeap []topIndex
    16  
    17  func (q topHeap) Len() int           { return len(q) }
    18  func (q topHeap) Less(i, j int) bool { return q[i].Count < q[j].Count }
    19  func (q topHeap) Swap(i, j int) {
    20  	q[i], q[j] = q[j], q[i]
    21  }
    22  
    23  func (q *topHeap) Push(x interface{}) {
    24  	t := x.(topIndex)
    25  	*q = append(*q, t)
    26  }
    27  
    28  func (q *topHeap) Pop() interface{} {
    29  	old := *q
    30  	n := len(old)
    31  	x := old[n-1]
    32  	*q = old[:n-1]
    33  	return x
    34  }
    35  
    36  func topn(counts []int, n int) []int {
    37  	if n == 0 {
    38  		log.Panicf("topn: n=0, counts=%d", counts)
    39  	}
    40  	if m := len(counts); m < n {
    41  		n = m
    42  	}
    43  	var q topHeap
    44  	for i, count := range counts {
    45  		if len(q) < n {
    46  			heap.Push(&q, topIndex{i, count})
    47  		} else if count > q[0].Count {
    48  			heap.Pop(&q)
    49  			heap.Push(&q, topIndex{i, count})
    50  		}
    51  	}
    52  	indices := make([]int, n)
    53  	for i := range indices {
    54  		indices[i] = q[i].Index
    55  	}
    56  	return indices
    57  }