cadmin/server/package/graveserver/leafstalk/otherutils/datastruct/smallheap/smallheap.go

// 最小堆存放数据项
// 并实现了一个时序队列，支持添加、删除、修改等操作
// 实现了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)
}