dht/closest_nodes.go

   1 package dht
   2
   3 import (
   4         "container/heap"
   5 )
   6
   7 type nodeMaxHeap struct {
   8         IDs    []nodeID
   9         Target nodeID
  10 }
  11
  12 func (mh nodeMaxHeap) Len() int { return len(mh.IDs) }
  13
  14 func (mh nodeMaxHeap) Less(i, j int) bool {
  15         m := mh.IDs[i].Distance(&mh.Target)
  16         n := mh.IDs[j].Distance(&mh.Target)
  17         return m.Cmp(&n) > 0
  18 }
  19
  20 func (mh *nodeMaxHeap) Pop() (ret interface{}) {
  21         ret, mh.IDs = mh.IDs[len(mh.IDs)-1], mh.IDs[:len(mh.IDs)-1]
  22         return
  23 }
  24 func (mh *nodeMaxHeap) Push(val interface{}) {
  25         mh.IDs = append(mh.IDs, val.(nodeID))
  26 }
  27 func (mh nodeMaxHeap) Swap(i, j int) {
  28         mh.IDs[i], mh.IDs[j] = mh.IDs[j], mh.IDs[i]
  29 }
  30
  31 type closestNodesSelector struct {
  32         closest nodeMaxHeap
  33         k       int
  34 }
  35
  36 func (cns *closestNodesSelector) Push(id nodeID) {
  37         heap.Push(&cns.closest, id)
  38         if cns.closest.Len() > cns.k {
  39                 heap.Pop(&cns.closest)
  40         }
  41 }
  42
  43 func (cns *closestNodesSelector) IDs() []nodeID {
  44         return cns.closest.IDs
  45 }
  46
  47 func newKClosestNodesSelector(k int, targetID nodeID) (ret closestNodesSelector) {
  48         ret.k = k
  49         ret.closest.Target = targetID
  50         return
  51 }