LeetCode: Partition List

Given a linked list and a value x, partition it such that all nodes less than x come before nodes greater than or equal to x.

You should preserve the original relative order of the nodes in each of the two partitions.

For example,
Given 1->4->3->2->5->2 and x = 3,
return 1->2->2->4->3->5.

---

The solution is trivial. We can set up two pointers to keep linking nodes that are less than the given value and linking nodes that are greater and equal to the given value. And at then end concatenate the two list together.

C++ Code:

/*
 * func: partition_list
 * goal: partition the list so that all nodes less thanx come before nodes greater
 *       than or equal to x.
 * @param head: head node of the list
 * @param x: target value
 * return: new head node
 */
/*
 * Set up a poiter to keep track of the nodes that greater and equal to the given value and another
 * pointer for nodes that less than the given value
 * complexity: time O(n), space O(1)
 */
ListNode *partition_list(ListNode *head, int x){
    ListNode *new_head = nullptr;
    ListNode *greater = nullptr;
    ListNode *greater_head = nullptr;
    ListNode *less = nullptr;
    ListNode *less_head = nullptr;
    ListNode *iter = head;
    while(iter != nullptr){
        if(iter->val >= x){
            if(!greater){
                greater = iter;
                greater_head = iter;
            }else{
                greater->next = iter;
                greater = greater->next;
            }
            iter = iter->next;
        }else{
            if(!less){
                less = iter;
                less_head = iter;
            }else{
                less->next = iter;
                less = less->next;
            }
            iter = iter->next;
        }
    }
    
    //Set the head of less list to the new head if less list exists
    if(less){
        less->next = greater_head;
        new_head = less_head;
    }else{
        new_head = greater_head;
    }
    //Set the tail node's next to null
    if(greater){
        greater->next = nullptr;
    }
    
    return new_head;

}

Python Code:

# func: rearrange the list so that nodes less than the given value to the left and others to the right
# @param head: head node of the list
# @param x: the given value
# @return: the new head of the list
# complexity: time O(n), space O(1)
def partition_list(head, x):
    iter = head
    less_iter = None
    less_head = None
    greater_iter = None
    greater_head = None

    while iter:
        if iter.val >= x:
            if not greater_head:
                greater_head = iter
                greater_iter = iter
            else:
                greater_iter.next = iter
                greater_iter = greater_iter.next
        else:
            if not less_head:
                less_head = iter
                less_iter = iter
            else:
                less_iter.next = iter
                less_iter = less_iter.next

        iter = iter.next

    if not less_head:
        new_head = greater_head
    else:
        less_iter.next = greater_head
        new_head = less_head

    if greater_head:
        greater_iter.next = None

    return new_head