advent/src/advent/heaphash.cr

66 lines
1.1 KiB
Crystal

require "./heap.cr"
# Inspired by Python's heapdict
class HeapHash(K,V)
private class Wrapper(K,V)
property key : K
property value : V
property index : Int32
def initialize(@key, @value, @index)
end
end
private class InternalHeap(K, V) < Array(Wrapper(K,V))
def swap(i,j)
super(i,j)
self[i].index = i
self[j].index = j
end
end
def initialize
@hash = {} of K => Wrapper(K,V)
@heap = InternalHeap(K,V).new
end
def delete(k)
wrapper = @hash[k]
@heap.bubble_up(wrapper.index) do |i, j|
false
end
pop
end
def []=(k : K, v : V)
if @hash.has_key? k
delete k
end
wrapper = Wrapper(K,V).new(k, v, @heap.size)
@hash[k] = wrapper
@heap.heap_push(wrapper) do |i, j|
i.value < j.value
end
end
def [](k)
wrapper = @hash[k]
wrapper.value
end
def []?(k)
return nil unless wrapper = @hash[k]?
wrapper.value
end
def pop
wrapper = @heap.heap_pop do |l, r|
l.value < r.value
end
@hash.delete wrapper.key
{wrapper.key, wrapper.value}
end
delegate size, empty?, to: @heap
end