go-hep.org/x/hep@v0.38.1/fastjet/internal/heap/heap.go (about)

     1  // Copyright ©2017 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package heap
     6  
     7  // pair is an item in the heap. It contains the indices of the two jets to be
     8  // clustered and their kt distance.
     9  type pair struct {
    10  	// jeti and jetj are the indices of the jets to be recombined.
    11  	jeti, jetj int
    12  	// dij is the kt distance for the two particles.
    13  	dij float64
    14  }
    15  
    16  // Heap contains a slice of items and the last index.
    17  type Heap struct {
    18  	n     int // n is the last index of the slice
    19  	items []pair
    20  }
    21  
    22  // New returns a heap pointer.
    23  func New() *Heap {
    24  	h := &Heap{n: 0}
    25  	h.items = make([]pair, 1)
    26  	return h
    27  }
    28  
    29  // Push inserts two new clustering candidates and their kt distance.
    30  func (h *Heap) Push(jeti, jetj int, dij float64) {
    31  	item := pair{jeti: jeti, jetj: jetj, dij: dij}
    32  	h.n++
    33  	if h.n >= len(h.items) {
    34  		h.items = append(h.items, item)
    35  	} else {
    36  		h.items[h.n] = item
    37  	}
    38  	h.moveUp(h.n)
    39  }
    40  
    41  // Pop returns the two jets with the smallest distance.
    42  // It returns -1, -1, 0 if the heap is empty.
    43  func (h *Heap) Pop() (jeti, jetj int, dij float64) {
    44  	if h.n == 0 {
    45  		return -1, -1, 0
    46  	}
    47  	item := h.items[1]
    48  	h.n--
    49  	if h.n == 0 {
    50  		return item.jeti, item.jetj, item.dij
    51  	}
    52  	h.swap(1, h.n+1)
    53  	h.moveDown(1)
    54  	return item.jeti, item.jetj, item.dij
    55  }
    56  
    57  // IsEmpty returns whether a heap is empty.
    58  func (h *Heap) IsEmpty() bool {
    59  	return h.n == 0
    60  }
    61  
    62  // moveUp compares an item with its parent and moves it up if it has
    63  // a smaller distance. It will keep moving up until the parent's distance is smaller
    64  // or if it reaches the top.
    65  func (h *Heap) moveUp(i int) {
    66  	for {
    67  		parent := int(i / 2)
    68  		if parent == 0 {
    69  			return
    70  		}
    71  		if h.less(i, parent) {
    72  			h.swap(i, parent)
    73  			i = parent
    74  			continue
    75  		}
    76  		return
    77  	}
    78  }
    79  
    80  // moveDown compares an item to its children and moves it down
    81  // if one of the children has a smaller distance. It will keep
    82  // moving down until it reaches a leaf or both children have
    83  // a bigger distance.
    84  func (h *Heap) moveDown(i int) {
    85  	for {
    86  		left := 2 * i
    87  		if left > h.n {
    88  			return
    89  		}
    90  		right := 2*i + 1
    91  		var smallestChild int
    92  		switch {
    93  		case left == h.n:
    94  			smallestChild = left
    95  		case h.less(left, right):
    96  			smallestChild = left
    97  		default:
    98  			smallestChild = right
    99  		}
   100  		if h.less(smallestChild, i) {
   101  			h.swap(i, smallestChild)
   102  			i = smallestChild
   103  			continue
   104  		}
   105  		return
   106  	}
   107  }
   108  
   109  // less returns whether item at index i has a smaller distance
   110  // than the item at index j.
   111  func (h *Heap) less(i, j int) bool {
   112  	return h.items[i].dij < h.items[j].dij
   113  }
   114  
   115  // swap swaps the two items at the indices i and j.
   116  func (h *Heap) swap(i, j int) {
   117  	h.items[i], h.items[j] = h.items[j], h.items[i]
   118  }