sharparam.com
Code for the Advent of Code event
aoc
advent-of-code
1# frozen_string_literal: true
2
3class PriorityQueue
4 include Enumerable
5
6 class Element
7 include Comparable
8
9 attr_accessor :value, :priority
10
11 def initialize(value, priority)
12 @value, @priority = value, priority
13 end
14
15 def <=>(other)
16 @priority <=> other.priority
17 end
18 end
19
20 def initialize
21 @elements = [nil]
22 end
23
24 def size
25 @elements.size - 1
26 end
27
28 def each(&)
29 @elements.drop(1).each(&)
30 end
31
32 def empty?
33 size == 0
34 end
35
36 def <<(element)
37 @elements << element
38 bubble_up(@elements.size - 1)
39 end
40
41 def enqueue(value, priority = nil)
42 value = Element.new(value, priority) if priority
43 self << value
44 end
45
46 def dequeue
47 exchange(1, @elements.size - 1)
48 element = @elements.pop.tap { bubble_down(1) }
49 return element.value if element.is_a? Element
50 element
51 end
52
53 private
54
55 def bubble_up(index)
56 parent_index = index / 2
57 return if index <= 1
58 return if @elements[parent_index] <= @elements[index]
59 exchange(index, parent_index)
60 bubble_up(parent_index)
61 end
62
63 def bubble_down(index)
64 child_index = index * 2
65 max_index = @elements.size - 1
66 return if child_index > max_index
67 not_last_element = child_index < max_index
68 left_element = @elements[child_index]
69 right_element = @elements[child_index + 1]
70 child_index += 1 if not_last_element && right_element < left_element
71 return if @elements[index] <= @elements[child_index]
72 exchange(index, child_index)
73 bubble_down(child_index)
74 end
75
76 def exchange(source, target)
77 @elements[source], @elements[target] = @elements[target], @elements[source]
78 end
79end