admin-entrance/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
	GetSortKey() int64
	GetUniqueKey() int64
}

type SmallHeap[T SortItem] struct {
	Items      []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 SmallHeap[T]) Len() int {
	return len(h.Items)
}

func (h SmallHeap[T]) Less(i, j int) bool {
	return h.Items[i].GetSortKey() < h.Items[j].GetSortKey()
}

func (h SmallHeap[T]) Swap(i, j int) {
	h.Items[i], h.Items[j] = h.Items[j], h.Items[i]
	h.Items[i].SetIndex(i)
	h.Items[j].SetIndex(j)
}

func (h *SmallHeap[T]) Push(x interface{}) {
	v, ok := x.(T)
	if ok {
		v.SetIndex(len(h.Items))
		h.Items = append(h.Items, v)
	}
}

func (h *SmallHeap[T]) Pop() interface{} {
	n := len(h.Items)
	x := h.Items[n-1]
	h.Items = h.Items[:n-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)

	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, it)
	k := it.GetUniqueKey()
	h.QueryTable[k] = it
	return nil
}

// 删除
func (h *SmallHeap[T]) Remove(index int) T {
	if index < 0 || index >= h.Len() {
		var t T
		return t
	}

	it := heap.Remove(h, index)
	v, ok := it.(T)
	if ok {
		v.SetIndex(-1)
		delete(h.QueryTable, v.GetUniqueKey())
		return v
	}
	var t T
	return t
}

func (h *SmallHeap[T]) RemoveByUniqueKey(uniqueKey int64) T {
	v, ok := h.QueryTable[uniqueKey]
	if !ok {
		var t T
		return t
	}
	return h.Remove(v.GetIndex())
}

// 重新排序
func (h *SmallHeap[T]) Resort(t1 T) bool {
	old, ok := h.GetByUniqueKey(t1.GetUniqueKey())
	if ok {
		heap.Fix(h, old.GetIndex())
		return true
	}
	return false
}

// 查看切片中的元素
func (h *SmallHeap[T]) Get(pos int) (T, error) {
	var zero T
	if pos < 0 || pos >= h.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)
}