Partition List Python Solution with Tests

lcpy gen -n 86   # by problem number
lcpy gen -s partition_list   # by problem name

Input: head = [1,4,3,2,5,2], x = 3
Output: [1,2,2,4,3,5]

Input: head = [2,1], x = 2
Output: [1,2]

from leetcode_py import ListNode


class Solution:
    # Time: O(n)
    # Space: O(1)
    def partition(self, head: ListNode[int] | None, x: int) -> ListNode[int] | None:
        before_head = before = ListNode[int](0)
        after_head = after = ListNode[int](0)
        current = head
        while current:
            if current.val < x:
                before.next = current
                before = before.next
            else:
                after.next = current
                after = after.next
            current = current.next
        after.next = None
        before.next = after_head.next
        return before_head.next