Go语言通用优先级队列：基于接口的实现与container/heap对比分析

// prio.Interface 定义了可插入优先级队列的元素的行为。
type Interface interface {
    // Less 返回当前元素是否应排在元素x之前（优先级更高）。
    Less(x Interface) bool
    // Index 由优先级队列调用，当此元素移动到索引i时更新其位置。
    Index(i int)
}

// Queue 表示一个优先级队列。
// Queue的零值是一个空的队列，可以直接使用。
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 返回一个用给定元素初始化的优先级队列。
// 形式为 New(x...) 的调用使用x的底层数组来实现队列，因此可能会改变x的元素。
// 复杂度为 O(n)，其中 n = len(x)。
func New(x ...Interface) Queue {
    q := Queue{x}
    heapify(q.h)
    return q
}

// Push 将元素x推入队列。
// 复杂度为 O(log(n))，其中 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) 由 up 完成。
}

// Pop 从队列中移除一个最小元素（根据Less方法），并返回它。
// 复杂度为 O(log(n))，其中 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) 由 down 完成。
    }
    q.h = h
    x.Index(-1) // 为了安全，将移除元素的索引设为-1
    return x
}

// Peek 返回队列中的最小元素（根据Less方法），但不移除它。
func (q *Queue) Peek() Interface {
    if len(q.h) == 0 {
        return nil // 或者返回错误，取决于具体需求
    }
    return q.h[0]
}

// Remove 移除队列中索引为i的元素并返回它。
// 复杂度为 O(log(n))，其中 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 // 将最后一个元素移到指定位置i，并清除原位置
    h = h[:n]              // 缩短切片
    if i < n {             // 如果移除的不是最后一个元素
        down(h, i) // h[i].Index(i) 由 down 完成。
        up(h, i)   // 元素可能需要向上或向下调整
    }
    q.h = h
    x.Index(-1) // 为了安全，将移除元素的索引设为-1
    return x
}

// Len 返回队列中的元素数量。
func (q *Queue) Len() int {
    return len(q.h)
}

// heapify 在 O(n) 时间内建立堆不变性。
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) 由 down 完成。
        down(h, i)
    }
}

// up 将位置i的元素向上移动以恢复堆不变性。
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                        // 继续向上检查
    }
}

// down 将位置i的元素向下移动以恢复堆不变性。
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下
)

// MyInt 是一个简单的整数类型，实现 prio.Interface
type MyInt int

func (x MyInt) Less(y prio.Interface) bool {
    return x < y.(MyInt) // 比较整数值
}

func (x MyInt) Index(i int) {
    // 对于简单的整数，如果不需要Remove，可以不存储索引
    // 但为了完整性，这里可以留空或打印
    // fmt.Printf("MyInt %d moved to index %d\n", x, i)
}

func main() {
    fmt.Println("--- 简单整数优先级队列 ---")
    pq := prio.New() // 创建一个空队列

    pq.Push(MyInt(5))
    pq.Push(MyInt(2))
    pq.Push(MyInt(8))
    pq.Push(MyInt(1))

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

    for pq.Len() > 0 {
        val := pq.Pop()
        fmt.Printf("Pop: %v\n", val)
    }
    // 预期输出: 1, 2, 5, 8
}

package main

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

// Task 表示一个任务，包含优先级和在队列中的索引
type Task struct {
    Name     string
    Priority int // 优先级值，越小优先级越高
    index    int // 在堆中的索引，用于快速移除
}

// Less 实现了 prio.Interface 的 Less 方法
func (t *Task) Less(other prio.Interface) bool {
    return t.Priority < other.(*Task).Priority // 优先级值越小，优先级越高
}

// Index 实现了 prio.Interface 的 Index 方法
func (t *Task) Index(i int) {
    t.index = i // 更新任务在堆中的索引
}

func main() {
    fmt.Println("\n--- 带有索引管理任务的优先级队列 ---")
    pq := prio.New()

    // 创建任务并推入队列
    taskA := &Task{Name: "Task A", Priority: 3}
    taskB := &Task{Name: "Task B", Priority: 1}
    taskC := &Task{Name: "Task C", Priority: 5}
    taskD := &Task{Name: "Task D", Priority: 2}

    pq.Push(taskA)
    pq.Push(taskB)
    pq.Push(taskC)
    pq.Push(taskD)

    fmt.Printf("队列长度: %d\n", pq.Len()) // 输出: 4
    fmt.Printf("Peek: %+v (Priority: %d)\n", pq.Peek().(*Task).Name, pq.Peek().(*Task).Priority) // 预期输出: Task B (Priority: 1)

    // 移除Task D (优先级为2)
    fmt.Printf("尝试移除 Task D (当前索引: %d)\n", taskD.index)
    removedTask := pq.Remove(taskD.index).(*Task) // 使用taskD.index来移除
    fmt.Printf("移除的元素: %+v (Priority: %d)\n", removedTask.Name, removedTask.Priority)
    fmt.Printf("移除后队列长度: %d\n", pq.Len()) // 输出: 3

    fmt.Println("移除后队列元素 (按优先级):")
    for pq.Len() > 0 {
        val := pq.Pop().(*Task)
        fmt.Printf("Pop: %+v (Priority: %d)\n", val.Name, val.Priority)
    }
    // 预期输出: Task B (Priority: 1), Task A (Priority: 3), Task C (Priority: 5)
}