LeetCode "21. Merge Two Sorted Lists"

merge_sorted_linkedlist.py

"""
>>> to_maybe_list = lambda x: [] if x is None else list(n.val for n in x)

>>> m = lambda a,b: to_maybe_list(merge(*map(ListNode.from_iter, (a, b))))
>>> assert(m([], []) == [])
>>> assert(m([], [0]) == [0])
>>> assert(m([0], []) == [0])
>>> assert(m([1,2,4], [1,3,4]) == [1,1,2,3,4,4])
>>> assert(m([1,2,5], [2,3,4]) == [1,2,2,3,4,5])
>>> assert(m([2,3,4], [1,2,5]) == [1,2,2,3,4,5])
>>> assert(m([1,1,1,2,3], [1,2,2,4,4]) == [1,1,1,1,2,2,2,3,4,4])

>>> to_maybe_list = lambda x: [] if x is None else list(n.val for n in x)
>>> m = lambda a,b: to_maybe_list(pick(*map(ListNode.from_iter, (a, b))))
>>> assert(m([], []) == [])
>>> assert(m([], [0]) == [0])
>>> assert(m([0], []) == [0])
>>> assert(m([1,2,4], [1,3,4]) == [1,1,2,3,4,4])
>>> assert(m([1,2,5], [2,3,4]) == [1,2,2,3,4,5])
>>> assert(m([2,3,4], [1,2,5]) == [1,2,2,3,4,5])
>>> assert(m([1,1,1,2,3], [1,2,2,4,4]) == [1,1,1,1,2,2,2,3,4,4])
"""

class LinkedList:
    """
    >>> empty = ListNode.from_iter([])
    >>> empty
    >>> list3= ListNode.from_iter([1,2,3])
    >>> list3
    ListNode(1, ...)
    >>> str(list3)
    '[1, 2, 3]'
    >>> iter3 = ListNode.from_iter(i for i in range(3))
    >>> str(iter3)
    '[0, 1, 2]'
    """
    def from_iter(I):
        head = n = None
        for v in I:
            if n is None:
                head = n = ListNode(v)
            else:
                n.next = n = ListNode(v)
        return head

    def __iter__(self):
        yield self
        if self.next:
            yield from self.next

    def __str__(self):
        return str([v.val for v in self])

    def __repr__(self):
        if self.next:
            return self.__class__.__name__ + f"({self.val}, ...)"
        else:
            return self.__class__.__name__ + f"({self.val})"

class ListNode(LinkedList):
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

def merge(A, B) -> "None|ListNode":
    "Merge B to A"
    if A is None or B is None: return A or B
    if A.val > B.val: A, B = B, A
    ans = A
    while A.next:
        if A.next.val > B.val:
            A.next, B = B, A.next
        A = A.next
    A.next = B
    return ans
        
def pick(A, B) -> "None|ListNode":
    "Pick one from A or B each iteration"
    if A is None or B is None: return A or B
    if A.val > B.val: A, B = B, A
    A.next = pick(A.next, B)
    return A