leetcode-go/leetcode/hard/23.go

package main

// MinPQ ...
type MinPQ []*ListNode

// Insert ...
func (pq *MinPQ) Insert(node *ListNode) {
	*pq = append(*pq, node)
	pq.swim(pq.Size())
}

// Min ...
func (pq *MinPQ) Min() *ListNode {
	return (*pq)[0]
}

// DelMin ...
func (pq *MinPQ) DelMin() *ListNode {
	min := (*pq)[0]
	size := pq.Size()
	pq.exch(1, size)
	*pq = (*pq)[:size-1]
	pq.sink(1)
	return min
}

// Empty ...
func (pq *MinPQ) Empty() bool {
	return pq.Size() == 0
}

// Size ...
func (pq *MinPQ) Size() int {
	return len(*pq)
}

// Note: i, j is NOT the index of element, but the rank of it.
func (pq *MinPQ) less(i, j int) bool {
	return (*pq)[i-1].Val < (*pq)[j-1].Val
}

func (pq *MinPQ) exch(i, j int) {
	(*pq)[i-1], (*pq)[j-1] = (*pq)[j-1], (*pq)[i-1]
}

func (pq *MinPQ) swim(k int) {
	for k > 1 && pq.less(k, k/2) {
		pq.exch(k, k/2)
		k /= 2
	}
}

func (pq *MinPQ) sink(k int) {
	size := pq.Size()
	for 2*k <= size { // !!
		j := 2 * k
		if j < size && pq.less(j+1, j) {
			j++ // Compare with the less child
		}
		if !pq.less(j, k) {
			break // If j !< k, ie. k <= j, break
		}
		pq.exch(k, j)
		k = j
	}
}

/**
 * Definition for singly-linked list.
 * type ListNode struct {
 *     Val int
 *     Next *ListNode
 * }
 */
func mergeKLists(lists []*ListNode) *ListNode {
	k := len(lists)
	if k == 0 {
		return nil
	}
	dummy := ListNode{}
	curr := &dummy
	var pq MinPQ
	for _, head := range lists {
		if head != nil {
			pq.Insert(head)
		}
	}
	for !pq.Empty() {
		min := pq.DelMin()
		curr.Next = min
		curr = curr.Next
		// pq is empty, or not empty && next min != min
		// Not empty && next min == min
		if pq.Empty() || pq.Min().Val != min.Val {
			if min.Next != nil {
				pq.Insert(min.Next)
			}
		} else {
			nexts := make([]*ListNode, 0, pq.Size())
			if min.Next != nil {
				nexts = append(nexts, min.Next)
			}
			for !pq.Empty() && pq.Min().Val == min.Val {
				min = pq.DelMin()
				curr.Next = min
				curr = curr.Next
				if min.Next != nil {
					nexts = append(nexts, min.Next)
				}
			}
			for _, next := range nexts {
				pq.Insert(next)
			}
		}
	}
	return dummy.Next
}

// func main() {
// 	l1 := toLinkedList([]int{-9, -7, -7})
// 	l2 := toLinkedList([]int{-6, -4})
// 	l3 := toLinkedList([]int{-6, -5})
// 	l4 := toLinkedList([]int{-9, -8, -6})
// 	l5 := toLinkedList([]int{-10})
// 	l6 := toLinkedList([]int{-5})
// 	lists := []*ListNode{l1, l2, l3, l4, l5, l6}
// 	var pq MinPQ
// 	pq.Insert(&ListNode{-9, nil})
// 	pq.Insert(&ListNode{-6, nil})
// 	pq.Insert(&ListNode{-6, nil})
// 	pq.Insert(&ListNode{-9, nil})
// 	pq.Insert(&ListNode{-10, nil})
// 	pq.Insert(&ListNode{-5, nil})
// 	for !pq.Empty() {
// 		println(pq.DelMin().Val)
// 	}
// 	printList(mergeKLists(lists))
// }