725.Split Linked List in Parts

/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode(int x) { val = x; }
 * }
 */
class Solution {
    public ListNode[] splitListToParts(ListNode root, int k) {
        Queue<Integer> q = new LinkedList<>();

        // count size
        ListNode temp = root;
        int size = 0;
        while (temp != null){
            size++;
            temp = temp.next;
        }

        int base = size / k;
        int remain = size % k;
        // k slots, if remain > 0, add 1 into each slot
        // 10 / 3 = 3, remain = 1, k = 3, then 4, 3, 3
        for (int i = 1; i <= k; i++){
            if (remain > 0){
                q.offer(base + 1);
                remain--;
            }
            else{
                q.offer(base);
            }
        }

        System.out.println(q);

        ListNode[] ans = new ListNode[k];
        int count = q.poll();
        int index = 0;
        ans[index] = root;
        while (root != null && !q.isEmpty()){
            if (count == 0){
                count = q.poll();
                index++;
                ans[index] = root;
            }

            // move to next, count--
            count--;
            ListNode nextNode = root.next;
            if (count == 0) root.next = null;
            root = nextNode;
        }

        return ans;
    }
}