leetcode/linklist/merge_two_lists.py

from typing import Optional
from utils import ListNode, arr_to_linklist, linkList_to_arr
import pytest


class Solution:
    def mergeTwoLists(self, list1: Optional[ListNode], list2: Optional[ListNode]) -> Optional[ListNode]:
        # 递归
        if not list1:
            return list2
        elif not list2:
            return list1
        elif list1.val <= list2.val:
            list1.next = self.mergeTwoLists(list1.next, list2)
            return list1
        else:
            list2.next = self.mergeTwoLists(list2.next, list1)
            return list2


        # 非递归
        # list3 = ListNode()
        # pre = list3
        # while list1 and list2:
        #     if list1.val <= list2.val:
        #         pre.next = list1
        #         list1 = list1.next
        #     else:
        #         pre.next = list2
        #         list2 = list2.next
        #
        #     pre = pre.next
        # pre.next = list1 if list1 else list2
        # return list3.next

if __name__ == "__main__":
    s = Solution()
    list1 = arr_to_linklist([1, 2, 4])
    list2 = arr_to_linklist([1, 3, 4])
    res = s.mergeTwoLists(list1, list2)
    res = linkList_to_arr(res)
    assert res == [1,1,2,3,4,4]