feaf

pkg/merge_k_sorted_lists/main.go

@@ -0,0 +1,63 @@
+package mergeksortedlists
+
+import "container/heap"
+
+type ListNode struct {
+	Val  int
+	Next *ListNode
+}
+
+type ListNodeHeap []*ListNode
+
+func (l ListNodeHeap) Len() int {
+	return len(l)
+}
+
+func (l ListNodeHeap) Less(i int, j int) bool {
+	if l[i] == nil {
+		return true
+	}
+
+	if l[j] == nil {
+		return false
+	}
+
+	return l[i].Val < l[j].Val
+}
+
+func (l *ListNodeHeap) Pop() (popped any) {
+	*l, popped = (*l)[:len(*l)-1], (*l)[len(*l)-1]
+	return
+}
+
+func (l *ListNodeHeap) Push(x any) {
+	*l = append(*l, x.(*ListNode))
+}
+
+func (l ListNodeHeap) Swap(i int, j int) {
+	l[i], l[j] = l[j], l[i]
+}
+
+func mergeKLists(lists []*ListNode) *ListNode {
+	hp := ListNodeHeap(lists)
+	heap.Init(&hp)
+
+	sorted := ListNode{}
+	head := &sorted
+
+	for len(hp) > 0 {
+		min := heap.Pop(&hp).(*ListNode)
+		if min == nil {
+			continue
+		}
+
+		rest := min.Next
+
+		head.Next, head = min, min
+		if rest != nil {
+			heap.Push(&hp, rest)
+		}
+	}
+
+	return sorted.Next
+}