Spaces:
Running
Running
File size: 3,524 Bytes
b110593 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
// _ _
// __ _____ __ ___ ___ __ _| |_ ___
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
// \ V V / __/ (_| |\ V /| | (_| | || __/
// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
//
// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
//
// CONTACT: [email protected]
//
package priorityqueue
type supportedValueType interface {
any | uint64
}
// Item represents a queue item supporting an optional additional Value
type Item[T supportedValueType] struct {
ID uint64
Dist float32
Rescored bool
Value T
}
// Queue is a priority queue supporting generic item values
type Queue[T supportedValueType] struct {
items []Item[T]
less func(items []Item[T], i, j int) bool
}
// NewMin constructs a priority queue which prioritizes items with smaller distance
func NewMin[T supportedValueType](capacity int) *Queue[T] {
return &Queue[T]{
items: make([]Item[T], 0, capacity),
less: func(items []Item[T], i, j int) bool {
return items[i].Dist < items[j].Dist
},
}
}
// NewMax constructs a priority queue which prioritizes items with greater distance
func NewMax[T supportedValueType](capacity int) *Queue[T] {
return &Queue[T]{
items: make([]Item[T], 0, capacity),
less: func(items []Item[T], i, j int) bool {
return items[i].Dist > items[j].Dist
},
}
}
// Pop removes the next item in the queue and returns it
func (q *Queue[T]) Pop() Item[T] {
out := q.items[0]
q.items[0] = q.items[len(q.items)-1]
q.items = q.items[:len(q.items)-1]
q.heapify(0)
return out
}
// Top peeks at the next item in the queue
func (q *Queue[T]) Top() Item[T] {
return q.items[0]
}
// Len returns the length of the queue
func (q *Queue[T]) Len() int {
return len(q.items)
}
// Cap returns the remaining capacity of the queue
func (q *Queue[T]) Cap() int {
return cap(q.items)
}
// Reset clears all items from the queue
func (q *Queue[T]) Reset() {
q.items = q.items[:0]
}
// ResetCap drops existing queue items, and allocates a new queue with the given capacity
func (q *Queue[T]) ResetCap(capacity int) {
q.items = make([]Item[T], 0, capacity)
}
// Insert creates a valueless item and adds it to the queue
func (q *Queue[T]) Insert(id uint64, distance float32) int {
item := Item[T]{
ID: id,
Dist: distance,
}
return q.insert(item)
}
// InsertWithValue creates an item with a T type value and adds it to the queue
func (q *Queue[T]) InsertWithValue(id uint64, distance float32, val T) int {
item := Item[T]{
ID: id,
Dist: distance,
Value: val,
}
return q.insert(item)
}
func (q *Queue[T]) insert(item Item[T]) int {
q.items = append(q.items, item)
i := len(q.items) - 1
for i != 0 && q.less(q.items, i, q.parent(i)) {
q.swap(i, q.parent(i))
i = q.parent(i)
}
return i
}
func (q *Queue[T]) left(i int) int {
return 2*i + 1
}
func (q *Queue[T]) right(i int) int {
return 2*i + 2
}
func (q *Queue[T]) parent(i int) int {
return (i - 1) / 2
}
func (q *Queue[T]) swap(i, j int) {
q.items[i], q.items[j] = q.items[j], q.items[i]
}
func (q *Queue[T]) heapify(i int) {
left := q.left(i)
right := q.right(i)
smallest := i
if left < len(q.items) && q.less(q.items, left, i) {
smallest = left
}
if right < len(q.items) && q.less(q.items, right, smallest) {
smallest = right
}
if smallest != i {
q.swap(i, smallest)
q.heapify(smallest)
}
}
|