Go数据结构之HeapMap实现指定Key删除堆

            10
         /     \            
      15        30  
     /  \      /  \
   40    50  100   40

// Item represents an item in the priority queue.
type Item struct {
 value    int // 项的值。
 priority int // 项的优先级。
 index    int // 项在队列中的索引。
}

// PriorityQueue 实现了 heap.Interface 接口。
type PriorityQueue []*Item
// Len 返回优先级队列的长度。
func (pq PriorityQueue) Len() int {
     return len(pq)
}
// Less 比较优先级队列中的两个项。
func (pq PriorityQueue) Less(i, j int) bool {
     return pq[i].priority < pq[j].priority
}
// Swap 交换优先级队列中的两个项。
func (pq PriorityQueue) Swap(i, j int) {
     pq[i], pq[j] = pq[j], pq[i]
     pq[i].index = i
     pq[j].index = j
}
// Push 向优先级队列中添加一个项。
func (pq *PriorityQueue) Push(x interface{}) {
     item := x.(*Item)
     item.index = len(*pq)
     *pq = append(*pq, item)
}
// Pop 移除并返回优先级队列中的最小项。
func (pq *PriorityQueue) Pop() interface{} {
     old := *pq
     n := len(old)
     item := old[n-1]
     old[n-1] = nil // 避免内存泄漏
     item.index = -1
     *pq = old[0 : n-1]
     return item
}
// Update 修改项的优先级并更新其在优先级队列中的位置。
func (pq *PriorityQueue) Update(item *Item, value, priority int) {
 item.value = value
 item.priority = priority
 heap.Fix(pq, item.index)
}

// PriorityQueue facilitates queue of Items, providing Push, Pop, and
// PopByKey convenience methods. The ordering (priority) is an int64 value
// with the smallest value is the highest priority. PriorityQueue maintains both
// an internal slice for the queue as well as a map of the same items with their
// keys as the index. This enables users to find specific items by key. The map
// must be kept in sync with the data slice.
// See https://golang.org/pkg/container/heap/#example__priorityQueue
type PriorityQueue struct {
   // data is the internal structure that holds the queue, and is operated on by
   // heap functions
   data queue
   // dataMap represents all the items in the queue, with unique indexes, used
   // for finding specific items. dataMap is kept in sync with the data slice
   dataMap map[string]*Item
   // lock is a read/write mutex, and used to facilitate read/write locks on the
   // data and dataMap fields
   lock sync.RWMutex
}
// queue is the internal data structure used to satisfy heap.Interface. This
// prevents users from calling Pop and Push heap methods directly
type queue []*Item
// Item is something managed in the priority queue
type Item struct {
   // Key is a unique string used to identify items in the internal data map
   Key string
   // Value is an unspecified type that implementations can use to store
   // information
   Value interface{}
   // Priority determines ordering in the queue, with the lowest value being the
   // highest priority
   Priority int64
   // index is an internal value used by the heap package, and should not be
   // modified by any consumer of the priority queue
   index int
}

func (pq *PriorityQueue) Push(i *Item) error {
   if i == nil || i.Key == "" {
      return errors.New("error adding item: Item Key is required")
   }
   pq.lock.Lock()
   defer pq.lock.Unlock()
   if _, ok := pq.dataMap[i.Key]; ok {
      return ErrDuplicateItem
   }
   // Copy the item value(s) so that modifications to the source item does not
   // affect the item on the queue
   clone, err := copystructure.Copy(i)
   if err != nil {
      return err
   }
   pq.dataMap[i.Key] = clone.(*Item)
   heap.Push(&pq.data, clone)
   return nil
}

// PopByKey searches the queue for an item with the given key and removes it
// from the queue if found. Returns nil if not found. This method must fix the
// queue after removing any key.
func (pq *PriorityQueue) PopByKey(key string) (*Item, error) {
   pq.lock.Lock()
   defer pq.lock.Unlock()
   item, ok := pq.dataMap[key]
   if !ok {
      return nil, nil
   }
   // Remove the item the heap and delete it from the dataMap
   itemRaw := heap.Remove(&pq.data, item.index)
   delete(pq.dataMap, key)
   if itemRaw != nil {
      if i, ok := itemRaw.(*Item); ok {
         return i, nil
      }
   }
   return nil, nil
}

func main() {
    // 测试Push方法
    pq := &PriorityQueue{
        data:    queue{},
        dataMap: make(map[string]*Item),
    }
    item := &Item{
        Key:      "key1",
        Value:    "value1",
        Priority: 1,
        index:    0,
    }
    err := pq.Push(item)
    if err != nil {
        fmt.Println("Push error:", err)
    } else {
        fmt.Println("Push success")
    }
    // 测试Pop方法
    poppedItem, err := pq.Pop()
    if err != nil {
        fmt.Println("Pop error:", err)
    } else {
        fmt.Println("Popped item key:", poppedItem.Key)
    }
    // 测试PopByKey方法
    item2 := &Item{
        Key:      "key2",
        Value:    "value2",
        Priority: 2,
        index:    1,
    }
    pq.Push(item2)
    poppedItemByKey, err := pq.PopByKey("key2")
    if err != nil {
        fmt.Println("PopByKey error:", err)
    } else if poppedItemByKey == nil {
        fmt.Println("Item not found")
    } else {
        fmt.Println("Popped item by key:", poppedItemByKey.Key)
    }
    // 测试Len方法
    item3 := &Item{
        Key:      "key3",
        Value:    "value3",
        Priority: 3,
        index:    2,
    }
    pq.Push(item3)
    length := pq.Len()
    fmt.Println("Queue length:", length)
}