Split Linked List in Parts Python Solution

lcpy gen -n 725   # by problem number
lcpy gen -s split_linked_list_in_parts   # by problem name

Input: head = [1,2,3], k = 5
Output: [[1],[2],[3],[],[]]
Explanation: The first element output[0] has output[0].val = 1, output[0].next = null.
The last element output[4] is null, but its string representation as a ListNode is [].

Input: head = [1,2,3,4,5,6,7,8,9,10], k = 3
Output: [[1,2,3,4],[5,6,7],[8,9,10]]
Explanation: The input has been split into consecutive parts with size difference at most 1, and earlier parts are a larger size than the later parts.

from leetcode_py import ListNode


class Solution:
    # Time: O(n + k)
    # Space: O(k)
    def split_list_to_parts(self, head: ListNode[int] | None, k: int) -> list[ListNode[int] | None]:
        length = 0
        node = head
        while node is not None:
            length += 1
            node = node.next

        width, remainder = divmod(length, k)
        parts: list[ListNode[int] | None] = []
        node = head
        for i in range(k):
            parts.append(node)
            if node is None:
                continue
            part_size = width + (1 if i < remainder else 0)
            for _ in range(part_size - 1):
                if node.next is not None:
                    node = node.next
            next_head = node.next
            node.next = None
            node = next_head
        return parts