topk.go

Documentation: github.com/beorn7/perks/topk

     1  package topk
     2  
     3  import (
     4  	"sort"
     5  )
     6  
     7  // http://www.cs.ucsb.edu/research/tech_reports/reports/2005-23.pdf
     8  
     9  type Element struct {
    10  	Value string
    11  	Count int
    12  }
    13  
    14  type Samples []*Element
    15  
    16  func (sm Samples) Len() int {
    17  	return len(sm)
    18  }
    19  
    20  func (sm Samples) Less(i, j int) bool {
    21  	return sm[i].Count < sm[j].Count
    22  }
    23  
    24  func (sm Samples) Swap(i, j int) {
    25  	sm[i], sm[j] = sm[j], sm[i]
    26  }
    27  
    28  type Stream struct {
    29  	k   int
    30  	mon map[string]*Element
    31  
    32  	// the minimum Element
    33  	min *Element
    34  }
    35  
    36  func New(k int) *Stream {
    37  	s := new(Stream)
    38  	s.k = k
    39  	s.mon = make(map[string]*Element)
    40  	s.min = &Element{}
    41  
    42  	// Track k+1 so that less frequenet items contended for that spot,
    43  	// resulting in k being more accurate.
    44  	return s
    45  }
    46  
    47  func (s *Stream) Insert(x string) {
    48  	s.insert(&Element{x, 1})
    49  }
    50  
    51  func (s *Stream) Merge(sm Samples) {
    52  	for _, e := range sm {
    53  		s.insert(e)
    54  	}
    55  }
    56  
    57  func (s *Stream) insert(in *Element) {
    58  	e := s.mon[in.Value]
    59  	if e != nil {
    60  		e.Count++
    61  	} else {
    62  		if len(s.mon) < s.k+1 {
    63  			e = &Element{in.Value, in.Count}
    64  			s.mon[in.Value] = e
    65  		} else {
    66  			e = s.min
    67  			delete(s.mon, e.Value)
    68  			e.Value = in.Value
    69  			e.Count += in.Count
    70  			s.min = e
    71  		}
    72  	}
    73  	if e.Count < s.min.Count {
    74  		s.min = e
    75  	}
    76  }
    77  
    78  func (s *Stream) Query() Samples {
    79  	var sm Samples
    80  	for _, e := range s.mon {
    81  		sm = append(sm, e)
    82  	}
    83  	sort.Sort(sort.Reverse(sm))
    84  
    85  	if len(sm) < s.k {
    86  		return sm
    87  	}
    88  
    89  	return sm[:s.k]
    90  }
    91
View as plain text