Solution for merge k lists leetcode

challenge.py

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
    def mergeKLists(self, lists):
        """
        :type lists: List[ListNode]
        :rtype: ListNode
        """
        a = 0
        if len(lists) == 0:
            return None
        for i in range(0, len(lists)):
            if lists[i] is not None:
                a = i
                break
        if len(lists) > 1:
            for i in range(a+1, len(lists)):
                if lists[a] is None or lists[i] is None:
                    self.mergeTwoLists(lists[a], lists[i])
                    continue
                
                if lists[a].val > lists[i].val:
                    tmp = a
                    a = i
                    self.mergeTwoLists(lists[a], lists[tmp])
                else:
                    self.mergeTwoLists(lists[a], lists[i])

        return lists[a]

    def mergeTwoLists(self, list1, list2):
        old_head = list1
        prev = None
        while list2 is not None:
            while list1 is not None:
                if list1.next is None:
                    if list1.val <= list2.val:
                        tmp_next = list2.next
                        list1.next = list2
                        prev = list1
                        list1 = list1.next
                        list1.next = None
                        list2 = tmp_next
                    else:
                        if prev is None:
                            tmp_next = list2.next

                            list1, list1.next = list2, list1
                            old_head = list1

                            list2 = tmp_next
                    break
                elif list1.val <= list2.val <= list1.next.val:
                    list1.next, list2.next, list2 = \
                    list2,      list1.next, list2.next
                    prev = list1
                    list1 = list1.next
                    break
                else:
                    prev = list1
                    list1 = list1.next

        list1 = old_head