013.合并 K 个升序链表

• 23.合并 K 个升序链表

合并 K 个升序链表 - 小根堆

/**
 * Definition for singly-linked list.
 * type ListNode struct {
 *     Val int
 *     Next *ListNode
 * }
 */
func mergeKLists(lists []*ListNode) *ListNode {
	heap := NewMinHeap(len(lists))
	for _, list := range lists {
		if list != nil {
			heap.Push(list)
		}
	}
 
	dummy := &ListNode{}
	prev := dummy
	for heap.Size() > 0 {
		top := heap.Pop()
		if top.Next != nil {
			heap.Push(top.Next)
		}
		prev.Next = top
		prev = top
	}
 
	return dummy.Next
}
 
type MinHeap struct {
	tree []*ListNode
}
 
func NewMinHeap(capacity int) *MinHeap {
	return &MinHeap{
		tree: make([]*ListNode, 0, capacity),
	}
}
 
func (h *MinHeap) Size() int {
	return len(h.tree)
}
 
func (h *MinHeap) Push(val *ListNode) {
	h.tree = append(h.tree, val)
	h.Up(h.Size() - 1)
}
 
func (h *MinHeap) Pop() *ListNode {
	n := h.Size()
	if n == 0 {
		panic("heap is empty")
	}
	res := h.tree[0]
	h.swap(0, n - 1)
	h.tree = h.tree[:n - 1]
	h.Down(0)
	return res
}
 
func (h *MinHeap) Up(idx int) {
	for {
		parentIdx := (idx - 1) >> 1
		if parentIdx < 0 {
			break
		}
		if h.tree[idx].Val >= h.tree[parentIdx].Val {
			break
		}
		h.swap(idx, parentIdx)
		idx = parentIdx
	}
}
 
func (h *MinHeap) Down(idx int) {
	n := h.Size()
	for {
		minIdx := idx
		leftChildIdx := (idx << 1) + 1
		if leftChildIdx < n && h.tree[leftChildIdx].Val < h.tree[minIdx].Val {
			minIdx = leftChildIdx
		}
		if leftChildIdx + 1 < n && h.tree[leftChildIdx + 1].Val < h.tree[minIdx].Val {
			minIdx = leftChildIdx + 1
		}
		if minIdx == idx {
			break
		}
		h.swap(idx, minIdx)
		idx = minIdx
	}
}
 
func (h *MinHeap) swap(a, b int) {
	h.tree[a], h.tree[b] = h.tree[b], h.tree[a]
}