// 最小堆存放数据项 // 并实现了一个时序队列,支持添加、删除、修改等操作 // 实现了heap.Interface接口,提供O(logn)的插入和删除操作 // 同时使用map提供O(1)的查找操作 // 修改时间:2024-12-26 package smallheap import ( "container/heap" "errors" ) type SortItem interface { SetIndex(index int) GetIndex() int GetSortValue() int64 SetSortValue(int64) GetUniqueKey() int64 } type HeapSlice[T SortItem] []T type SmallHeap[T SortItem] struct { Items HeapSlice[T] QueryTable map[int64]T // sequence *snowflake.Node //1MS生成400万个,不分节点 } var ( ErrItemExists = errors.New("item already exists") ErrItemNotFound = errors.New("item not found") ) func (h HeapSlice[T]) Len() int { return len(h) } func (h HeapSlice[T]) Less(i, j int) bool { return h[i].GetSortValue() < h[j].GetSortValue() } func (h HeapSlice[T]) Swap(i, j int) { h[i], h[j] = h[j], h[i] h[i].SetIndex(i) h[j].SetIndex(j) } func (h *HeapSlice[T]) Push(x interface{}) { v, ok := x.(T) if ok { v.SetIndex(len(*h)) *h = append(*h, v) } } func (h *HeapSlice[T]) Pop() interface{} { n := len(*h) x := (*h)[n-1] *h = (*h)[:n-1] x.SetIndex(-1) // 设置索引为 -1 return x } func NewSmallHeap[T SortItem]() *SmallHeap[T] { sh := new(SmallHeap[T]) sh.Items = make([]T, 0, 16) sh.QueryTable = make(map[int64]T) // seq, err := snowflake.NewMsNode(0, 0) // if err != nil { // panic("NewSmallHeap init sequence error") // return nil // } // sh.sequence = seq heap.Init(&sh.Items) return sh } func (h *SmallHeap[T]) Clear() { h.Items = make([]T, 0) h.QueryTable = make(map[int64]T) } // 使用接口 // 添加 func (h *SmallHeap[T]) Add(it T) error { // 检查元素是否已存在 if _, exists := h.QueryTable[it.GetUniqueKey()]; exists { return ErrItemExists } heap.Push(&h.Items, it) k := it.GetUniqueKey() h.QueryTable[k] = it return nil } // 删除 func (h *SmallHeap[T]) Remove(index int) (T, error) { if index < 0 || index >= h.Items.Len() { var t T return t, ErrItemNotFound } it := heap.Remove(&h.Items, index) v, ok := it.(T) if ok { v.SetIndex(-1) delete(h.QueryTable, v.GetUniqueKey()) return v, nil } var t T return t, ErrItemNotFound } func (h *SmallHeap[T]) RemoveByUniqueKey(uniqueKey int64) (T, error) { v, ok := h.QueryTable[uniqueKey] if !ok { var t T return t, ErrItemNotFound } return h.Remove(v.GetIndex()) } // 重新排序 func (h *SmallHeap[T]) UpdateSortVal(key int64, sortVal int64) bool { old, ok := h.GetByUniqueKey(key) if ok { old.SetSortValue(sortVal) heap.Fix(&h.Items, old.GetIndex()) return true } return false } // 查看切片中的元素 func (h *SmallHeap[T]) Get(pos int) (T, error) { var zero T if pos < 0 || pos >= h.Items.Len() { return zero, ErrItemNotFound } return h.Items[pos], nil } func (h *SmallHeap[T]) GetByUniqueKey(uniqueKey int64) (T, bool) { v, ok := h.QueryTable[uniqueKey] return v, ok } // Size 返回堆中元素数量 // 时间复杂度: O(1) func (h *SmallHeap[T]) Size() int { return len(h.Items) }