diff --git a/aoc_test.go b/aoc_test.go index 27b6bec..ae579c9 100644 --- a/aoc_test.go +++ b/aoc_test.go @@ -336,6 +336,7 @@ func TestFibHeap(t *testing.T) { is.Equal(v, &elem{5, 8}) m := aoc.FibHeap(less) + m.Insert(&elem{1, 99}) m.Insert(&elem{12, 9}) m.Insert(&elem{11, 10}) m.Insert(&elem{10, 11}) @@ -343,11 +344,34 @@ func TestFibHeap(t *testing.T) { pq.Merge(m) + v = pq.Find(func(t *elem) bool { + return (*t)[0] == 6 + }) + is.Equal(v, &elem{6, 7}) + + v = pq.Find(func(t *elem) bool { + return (*t)[0] == 12 + }) + is.Equal(v, &elem{12, 9}) + + v = pq.ExtractMin() + is.True(v != nil) + is.Equal(v, &elem{1, 99}) + + pq.DecreaseKey( + func(t *elem) bool { return t[0] == 12 }, + func(t *elem) { t[0] = 3 }, + ) + + v = pq.ExtractMin() + is.True(v != nil) + is.Equal(v, &elem{3, 9}) + var keys []int for !pq.IsEmpty() { v := pq.ExtractMin() fmt.Println(v) keys = append(keys, v[0]) } - is.Equal(keys, []int{6, 7, 8, 9, 10, 11, 12}) + is.Equal(keys, []int{6, 7, 8, 9, 10, 11}) } diff --git a/search.go b/search.go index 80f6fb9..c2dc67c 100644 --- a/search.go +++ b/search.go @@ -166,11 +166,13 @@ func FindPath[C integer, N comparable](g pather[C, N], start, end N) (C, []N, ma } type fibTree[T any] struct { - value *T + value *T parent *fibTree[T] - child []*fibTree[T] + child []*fibTree[T] + mark bool } +func (t *fibTree[T]) Value() *T { return t.value } func (t *fibTree[T]) addAtEnd(n *fibTree[T]) { n.parent = t t.child = append(t.child, n) @@ -231,7 +233,7 @@ func (h *fibHeap[T]) ExtractMin() *T { } func (h *fibHeap[T]) consolidate() { - aux := make([]*fibTree[T], bits.Len(h.count)) + aux := make([]*fibTree[T], bits.Len(h.count)+1) for _, x := range h.trees { order := len(x.child) @@ -264,9 +266,83 @@ func (h *fibHeap[T]) consolidate() { func (h *fibHeap[T]) Merge(a *fibHeap[T]) { h.trees = append(h.trees, a.trees...) h.count += a.count - h.consolidate() + if h.least == nil || a.least != nil && h.less(a.least.value, h.least.value) { + h.least = a.least + } } -// func (h *fibHeap[T]) Find(n *T) *fibTree[T] { - -// } \ No newline at end of file +func (h *fibHeap[T]) find(fn func(*T) bool) *fibTree[T] { + var st []*fibTree[T] + st = append(st, h.trees...) + var tr *fibTree[T] + + for len(st) > 0 { + tr, st = st[0], st[1:] + ro := *tr.value + if fn(&ro) { + break + } + st = append(st, tr.child...) + } + + return tr +} + +func (h *fibHeap[T]) Find(fn func(*T) bool) *T { + if needle := h.find(fn); needle != nil { + return needle.value + } + + return nil +} + +func (h *fibHeap[T]) DecreaseKey(find func(*T) bool, decrease func(*T)) { + needle := h.find(find) + if needle == nil { + return + } + decrease(needle.value) + + if h.less(needle.value, h.least.value) { + h.least = needle + } + + if parent := needle.parent; parent != nil { + if h.less(needle.value, parent.value) { + h.cut(needle) + h.cascadingCut(parent) + } + } +} + +func (h *fibHeap[T]) cut(x *fibTree[T]) { + parent := x.parent + for i := range parent.child { + pos := parent.child[i] + if pos == x { + parent.child[i] = parent.child[len(parent.child)-1] + parent.child = parent.child[:len(parent.child)-1] + break + } + } + + x.parent = nil + x.mark = false + h.trees = append(h.trees, x) + + if h.less(x.value, h.least.value) { + h.least = x + } +} + +func (h *fibHeap[T]) cascadingCut(y *fibTree[T]) { + if y.parent != nil { + if !y.mark { + y.mark = true + return + } + + h.cut(y) + h.cascadingCut(y.parent) + } +}