partitionList - for Western Tech Interview Prep Q1

partitionList.java

/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode(int x) { val = x; }
 * }
 */
public class Solution {
    public ListNode partition(ListNode head, int x) {
        if (head == null || head.next == null) {
            return head;
        }
        ArrayList<Integer> greaterThanX = new ArrayList<Integer>();
        ArrayList<Integer> lessThanX = new ArrayList<Integer>();
        
        ListNode curr = head;
        while (true) {
            int val = curr.val;
            if (val >= x) {
                greaterThanX.add(val);
            } else {
                lessThanX.add(val);
            }
            if (curr.next != null) {
                curr = curr.next;
            } else {
                break;
            }
        }
        
        int lessThanXAmount = lessThanX.size();
        int greaterThanXAmount = greaterThanX.size();
        ListNode origHead = head;
        int i = 0;
        while (i < lessThanXAmount) {
            head.val = lessThanX.get(i);
            head = head.next;
            i++;
        }
        
        int j = 0;
        while (j < greaterThanXAmount) {
            head.val = greaterThanX.get(j);
            head = head.next;
            j++;
        }
        
        return origHead;
    }
}