当前位置：网站首页>go优先级队列实现

go优先级队列实现

2022-08-07 14:05:00 【Johns】

背景

最近在leetcode刷题, 遇到一个合并K个升序链表的问题, 就是把K个有序链表合并成一个有序链表

想到了一个idea就是把K个链表的头放到一个优先级队列, 每次弹最小的出来, 并把弹出来的节点的Next节点push回队列, 不断循环就能解决这个问题, 奈何go没有像Java一样提供那么易用的API, 我们只能使用container/heap实现一个简单的优先级队列.

实现

package main

import (
	"container/heap"
	"fmt"
)

type ListNode struct {
	Val  int
	Next *ListNode
}

// 定义一个优先级队列
type PriorityQueue []*ListNode

func (p PriorityQueue) Len() int {
	return len(p)
}

// 小于(<)就是小顶堆, 大于(>)就是大顶堆
func (p PriorityQueue) Less(i, j int) bool {
	return p[i].Val < p[j].Val
}
func (p PriorityQueue) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
	return
}

func (p *PriorityQueue) Push(x interface{}) {
	*p = append(*p, x.(*ListNode))
}

func (p *PriorityQueue) Pop() interface{} {
	old := *p
	n := len(old)
	x := old[n-1]
	*p = old[0 : n-1]
	return x
}

func main() {
	p := &PriorityQueue{}
	heap.Init(p)
	node1 := &ListNode{Val: 4}
	node2 := &ListNode{Val: 3}
	node3 := &ListNode{Val: 2}
	node4 := &ListNode{Val: 6}
	node5 := &ListNode{Val: 1}
	heap.Push(p, node1)
	heap.Push(p, node2)
	heap.Push(p, node3)
	heap.Push(p, node4)
	heap.Push(p, node5)

	for p.Len() > 0 {
		node := heap.Pop(p)
		fmt.Print(node.(*ListNode).Val, " ")
	}
}

输出:

1 2 3 4 6

合并K个有序链表

有了上面定义好的优先级队列, 我们可以开始解决问题了

func mergeKLists(lists []*ListNode) *ListNode {
    if len(lists)==0{
        return nil
    }
	dummy := &ListNode{Val: -1}
	p1 := dummy
	// 使用一个优先级队列（小根堆）， 每次把堆顶元素（最小值）弹出来
	p := &PriorityQueue{}
	heap.Init(p)
	for i := 0; i < len(lists); i++ {
        if lists[i]!=nil{
            heap.Push(p, lists[i])
        }	
	}

	for p.Len() > 0 {
		v := heap.Pop(p).(*ListNode)
		p1.Next = v
        if v.Next!=nil{
            heap.Push(p, v.Next)
        }
		p1 = p1.Next
	}
	return dummy.Next

}

// 测试
func main() {
	node1 := &ListNode{Val: 3}
	node2 := &ListNode{Val: 4}
	node1.Next = node2

	node3 := &ListNode{Val: 1}
	node4 := &ListNode{Val: 2}
	node3.Next = node4
	node5 := &ListNode{Val: 6}
	node4.Next = node5

	v := mergeKLists([]*ListNode{node1, node3})
	fmt.Println(v)
}
原网站

版权声明
本文为[Johns]所创，转载请带上原文链接，感谢
https://cloud.tencent.com/developer/article/2067873

边栏推荐

猜你喜欢

随机推荐