leetcode-practice-python/0023__merge_k_sorted_lists.py at master · Philosoft/leetcode-practice-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
"""
LeetCode link https://leetcode.com/problems/merge-k-sorted-lists/description/

You are given an array of k linked-lists lists, each linked-list is sorted in ascending order.

Merge all the linked-lists into one sorted linked-list and return it.

## Example 1

Input: lists = [[1,4,5],[1,3,4],[2,6]]
Output: [1,1,2,3,4,4,5,6]
Explanation:
    The linked-lists are:
    [
      1->4->5,
      1->3->4,
      2->6
    ]
    merging them into one sorted list:
    1->1->2->3->4->4->5->6

## Example 2

Input: lists = []
Output: []

## Example 3

Input: lists = [[]]
Output: []

## Constraints

* k == lists.length
* 0 <= k <= 10^4
* 0 <= lists[i].length <= 500
* -10^4 <= lists[i][j] <= 10^4
* lists[i] is sorted in ascending order.
* The sum of lists[i].length will not exceed 10^4.
"""
import math
from typing import List, Optional
from unittest import TestCase

from lib.ListNode import ListNode, build_list, flatten_list


class Solution(TestCase):
    def test_example_1(self):
        lists = list(map(build_list, [[1, 4, 5], [1, 3, 4], [2, 6]]))
        expected = [1, 1, 2, 3, 4, 4, 5, 6]

        self.assertEqual(expected, flatten_list(self.mergeKLists(lists)))

    def test_example_2(self):
        self.assertEqual([], flatten_list(self.mergeKLists([])))

    def test_example_3(self):
        self.assertEqual(None, self.mergeKLists([None]))

    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        dummy = ListNode()
        head = dummy
        while True:
            min_index = -1
            min_val = math.inf
            for idx, l in enumerate(lists):
                if not l:
                    continue

                if lists[idx].val < min_val:
                    min_val = lists[idx].val
                    min_index = idx

            if min_index == -1:
                break

            head.next = lists[min_index]
            lists[min_index] = lists[min_index].next
            head = head.next

        return dummy.next