diff --git a/pkg/merge_k_sorted_lists/main.go b/pkg/merge_k_sorted_lists/main.go new file mode 100644 index 0000000..c24be27 --- /dev/null +++ b/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 +}