leetcode

题目

23. 合并 K 个升序链表
    给你一个链表数组，每个链表都已经按升序排列。

请你将所有链表合并到一个升序链表中，返回合并后的链表。

代码

使用优先级队列

package listnode

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

// Len 计算长度
func (t PriorityQueue) Len() int {
	return len(t)
}

// Less 比较值
func (t PriorityQueue) Less(i, j int) bool {
	return t[i].Val < t[j].Val
}

// Swap 交换值
func (t PriorityQueue) Swap(i, j int) {
	t[i], t[j] = t[j], t[i]
}

// Push 存入值
func (t *PriorityQueue) Push(x interface{}) {
	node := x.(*ListNode)
	*t = append(*t, node)
}

// Pop 弹出栈
func (t *PriorityQueue) Pop() interface{} {
	old := *t
	n := len(old)
	node := old[n-1]
	*t = old[0 : n-1]
	return node
}

package listnode

import "container/heap"

type ListNode struct {
	Val  int
	Next *ListNode
}

// 23. 合并 K 个升序链表
func mergeKLists(lists []*ListNode) *ListNode {
	if len(lists) == 0 {
		return nil
	}
	// 定义虚拟头结点
	dumpy := &ListNode{-1, nil}
	p := dumpy
	// 优先级队列
	pq := make(PriorityQueue, 0)
	heap.Init(&pq)

	// 将k个链表的节点加入堆中
	for _, head := range lists {
		if head != nil {
			heap.Push(&pq, head)
		}
	}

	for pq.Len() > 0 {
		node := heap.Pop(&pq).(*ListNode)
		p.Next = node
		if node.Next != nil {
			heap.Push(&pq, node.Next)
		}
		p = p.Next
	}
	return dumpy.Next
}