Go语言中基于元素接口的优先级队列实现与container/heap的对比分析

type Interface interface {
    // Less 返回此元素是否应在元素x之前排序。
    Less(x Interface) bool
    // Index 在此元素被移动到索引i时，由优先级队列调用。
    Index(i int)
}

type Queue struct {
    h []Interface
}

// Copyright 2012 Stefan Nilsson
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package prio provides a priority queue.
// The queue can hold elements that implement the two methods of prio.Interface.
package prio

/*
A type that implements prio.Interface can be inserted into a priority queue.

The simplest use case looks like this:

        type myInt int

        func (x myInt) Less(y prio.Interface) bool { return x < y.(myInt) }
        func (x myInt) Index(i int)                {}

To use the Remove method you need to keep track of the index of elements
in the heap, e.g. like this:

        type myType struct {
                value int
                index int // index in heap
        }

        func (x *myType) Less(y prio.Interface) bool { return x.value < y.(*myType).value }
        func (x *myType) Index(i int)                { x.index = i }
*/
type Interface interface {
        // Less returns whether this element should sort before element x.
        Less(x Interface) bool
        // Index is called by the priority queue when this element is moved to index i.
        Index(i int)
}

// Queue represents a priority queue.
// The zero value for Queue is an empty queue ready to use.
type Queue struct {
        h []Interface
}

// New returns an initialized priority queue with the given elements.
// A call of the form New(x...) uses the underlying array of x to implement
// the queue and hence might change the elements of x.
// The complexity is O(n), where n = len(x).
func New(x ...Interface) Queue {
        q := Queue{x}
        heapify(q.h)
        return q
}

// Push pushes the element x onto the queue.
// The complexity is O(log(n)) where n = q.Len().
func (q *Queue) Push(x Interface) {
        n := len(q.h)
        q.h = append(q.h, x)
        up(q.h, n) // x.Index(n) is done by up.
}

// Pop removes a minimum element (according to Less) from the queue and returns it.
// The complexity is O(log(n)), where n = q.Len().
func (q *Queue) Pop() Interface {
        h := q.h
        n := len(h) - 1
        x := h[0]
        h[0], h[n] = h[n], nil
        h = h[:n]
        if n > 0 {
                down(h, 0) // h[0].Index(0) is done by down.
        }
        q.h = h
        x.Index(-1) // for safety
        return x
}

// Peek returns, but does not remove, a minimum element (according to Less) of the queue.
func (q *Queue) Peek() Interface {
        return q.h[0]
}

// Remove removes the element at index i from the queue and returns it.
// The complexity is O(log(n)), where n = q.Len().
func (q *Queue) Remove(i int) Interface {
        h := q.h
        n := len(h) - 1
        x := h[i]
        h[i], h[n] = h[n], nil
        h = h[:n]
        if i < n {
                down(h, i) // h[i].Index(i) is done by down.
                up(h, i)
        }
        q.h = h
        x.Index(-1) // for safety
        return x
}

// Len returns the number of elements in the queue.
func (q *Queue) Len() int {
        return len(q.h)
}

// Establishes the heap invariant in O(n) time.
func heapify(h []Interface) {
        n := len(h)
        for i := n - 1; i >= n/2; i-- {
                h[i].Index(i)
        }
        for i := n/2 - 1; i >= 0; i-- { // h[i].Index(i) is done by down.
                down(h, i)
        }
}

// Moves element at position i towards top of heap to restore invariant.
func up(h []Interface, i int) {
        for {
                parent := (i - 1) / 2
                if i == 0 || h[parent].Less(h[i]) {
                        h[i].Index(i)
                        break
                }
                h[parent], h[i] = h[i], h[parent]
                h[i].Index(i)
                i = parent
        }
}

// Moves element at position i towards bottom of heap to restore invariant.
func down(h []Interface, i int) {
        for {
                n := len(h)
                left := 2*i + 1
                if left >= n {
                        h[i].Index(i)
                        break
                }
                j := left
                if right := left + 1; right < n && h[right].Less(h[left]) {
                        j = right
                }
                if h[i].Less(h[j]) {
                        h[i].Index(i)
                        break
                }
                h[i], h[j] = h[j], h[i]
                h[i].Index(i)
                i = j
        }
}

package main

import (
    "fmt"
    "prio" // 假设prio包在你的GOPATH中
)

// 定义一个需要优先级排序的结构体
type Item struct {
    value string
    priority int
    index int // 存储其在堆中的索引
}

// 实现 prio.Interface 的 Less 方法
func (x *Item) Less(y prio.Interface) bool {
    return x.priority < y.(*Item).priority
}

// 实现 prio.Interface 的 Index 方法
func (x *Item) Index(i int) {
    x.index = i
}

func main() {
    // 创建一些 Item 实例
    item1 := &Item{value: "任务A", priority: 3}
    item2 := &Item{value: "任务B", priority: 1}
    item3 := &Item{value: "任务C", priority: 2}

    // 初始化优先级队列
    pq := prio.New(item1, item2, item3)

    fmt.Printf("队列长度: %d\n", pq.Len()) // 输出: 队列长度: 3

    // 查看最小元素
    minItem := pq.Peek().(*Item)
    fmt.Printf("最小元素: %s (优先级: %d)\n", minItem.value, minItem.priority) // 输出: 最小元素: 任务B (优先级: 1)

    // 弹出最小元素
    poppedItem := pq.Pop().(*Item)
    fmt.Printf("弹出元素: %s (优先级: %d)\n", poppedItem.value, poppedItem.priority) // 输出: 弹出元素: 任务B (优先级: 1)
    fmt.Printf("队列长度: %d\n", pq.Len()) // 输出: 队列长度: 2

    // 再次查看最小元素
    minItem = pq.Peek().(*Item)
    fmt.Printf("当前最小元素: %s (优先级: %d)\n", minItem.value, minItem.priority) // 输出: 当前最小元素: 任务C (优先级: 2)

    // 演示Remove方法，需要先找到索引
    // 假设我们想移除 item1 (任务A)
    // 在实际应用中，你可能需要一个map来根据value找到Item的指针，然后用其index字段来调用Remove
    // 这里我们直接使用 item1.index (在Push或New时，Index方法已被调用更新)
    fmt.Printf("任务A的当前索引: %d\n", item1.index) // 此时 item1.index 可能是0或1，取决于堆结构

    // 注意：这里的item1.index是在pq初始化后，item1被heapify或up/down操作时更新的。
    // 如果你直接用New初始化，item1的index可能不是你期望的，因为它可能已经被移动。
    // 正确的Remove通常用于在外部修改了某个元素的优先级后，需要更新其在堆中的位置，
    // 或者通过某种方式获取到元素的当前索引。
    // 为了演示，我们假设 item1 仍在队列中，并且我们知道其当前索引。
    // 在本例中，item1在初始化的切片中是第一个，但经过heapify，其index可能改变。
    // Pop后，item1的index可能被改变。
    // 假设我们现在知道item1的index是0 (如果它在堆顶)，或者1 (如果它在第二个位置)
    // 这里我们直接使用 item1.index 来移除
    if item1.index != -1 { // 检查元素是否仍在队列中
        removedItem := pq.Remove(item1.index).(*Item)
        fmt.Printf("移除元素: %s (优先级: %d)\n", removedItem.value, removedItem.priority)
        fmt.Printf("队列长度: %d\n", pq.Len())
    }
}