Merge K sorted lists of element (string, integers, ...)

"""
Merges k sorted lists of strings into one sorted list.

solve(["abc"]))           # -> ['a', 'b', 'c']
solve(["abc", "def"]))    # -> ['a', 'b', 'c', 'd', 'e', 'f']
solve(["abc", "abe"]))    # -> ['a', 'a', 'b', 'b', 'c', 'e']
"""

### STEP 2 ###
class Node:
    def __init__(self, val, lst_idx, elem_idx):
        self.val = val
        self.lst = lst_idx
        self.inx = elem_idx

class Min_Heap:
    def __init__(self):
        self.heap = []

    def add(self, val, lst_idx, elem_idx):
        self.heap.append(Node(val, lst_idx, elem_idx))

        index = len(self.heap) - 1
        while index > 0 and self.heap[index].val < self.heap[(index - 1) // 2].val:
            parent = (index - 1) // 2
            self.heap[index], self.heap[parent] = self.heap[parent], self.heap[index]
            index = parent

    def pop(self):
        element = (self.heap[0].val, self.heap[0].lst, self.heap[0].inx)

        last_index = len(self.heap) - 1
        self.heap[0] = self.heap[last_index]
        self.heap.pop(-1)

        i = 0
        while True:
            left = 2 * i + 1; right = 2 * i + 2
            smallest = i

            if left < len(self.heap) and self.heap[left].val < self.heap[smallest].val:
                smallest = left
            if right < len(self.heap) and self.heap[right].val < self.heap[smallest].val:
                smallest = right
            if smallest == i:
                break

            self.heap[i], self.heap[smallest] = self.heap[smallest], self.heap[i]
            i = smallest

        return element

    def is_empty(self):
        return len(self.heap) == 0

### STEP 1 ###
def solve_custom_heap(lists):
    """
    Time Complexity:  O(N log k) -- N = no. of chars; k = len('lists')
    Space Complexity: O(N)       -- N = no. of chars
    """
    heap = Min_Heap()
    result = []

    # Initialize the heap with the first element from each list
    for i, lst in enumerate(lists):
        if lst:
            heap.add(lst[0], i, 0)

    # Extract the smallest item and add the next item from the same list to the heap
    while not heap.is_empty():
        val, list_idx, elem_idx = heap.pop()
        result.append(val)
        # If there is a next element in the same list, push it into the heap
        if elem_idx + 1 < len(lists[list_idx]):
            next_val = lists[list_idx][elem_idx + 1]
            heap.add(next_val, list_idx, elem_idx + 1)

    return result

def solve_built_in(lists):
    """
    Time Complexity:  O(N log k) -- N = no. of chars; k = len('lists')
    Space Complexity: O(N)       -- N = no. of chars
    """
    answers = []

    import heapq
    heap = []

    # Initialize the heap with the first element from each list
    for i, lst in enumerate(lists):
        if lst:
            # Push a tuple of (string value, list index, element index)
            heapq.heappush(heap, (lst[0], i, 0))

    # Extract the smallest item and add the next item from the same list to the heap
    while heap:

        val, list_idx, element_idx = heapq.heappop(heap)
        answers.append(val)

        # If there is a next element in the same list, push it into the heap
        if element_idx + 1 < len(lists[list_idx]):
            next_val = lists[list_idx][element_idx + 1]
            heapq.heappush(heap, (next_val, list_idx, element_idx + 1))

    return answers