CareerCup 2.4

2.4 Write code to partition a linked list around a value x, such that all nodes less than x come before all nodes greater than or equal to x.

struct ListNode {
    int val;
    ListNode *next;
    ListNode(int x) : val(x), next(NULL) {}
};

ListNode* partition(ListNode *head, int x) {
    ListNode *beforeStart = NULL;
    ListNode *beforeEnd = NULL;
    ListNode *afterStart = NULL;
    ListNode *afterEnd = NULL;
    ListNode *node = head;

    while (node) {
        ListNode *next = node->next;
        node->next = NULL;
        if (node->val < x) {
            if (beforeStart == NULL) {
                beforeStart = beforeEnd = node;
            } else {
                beforeEnd->next = node;
                beforeEnd = node;
            }
        } else {
            if (afterStart == NULL) {
                afterStart = afterEnd = node;
            } else {
                afterEnd->next = node;
                afterEnd = node;
            }
        }
        node = next;
    }

    if (beforeStart == NULL) {
        return afterStart;
    } else {
        beforeEnd->next = afterStart;
        return beforeStart;
    }
}

Improved:

ListNode* partition(ListNode *head, int x) {
    ListNode *beforeStart = NULL;
    ListNode *afterStart = NULL;
    ListNode *node = head;

    while (node) {
        ListNode *next = node->next;
        if (node->val < x) {
            node->next = beforeStart;
            beforeStart = node;
        } else {
            node->next = afterStart;
            afterStart = node;
        }
        node = next;
    }

    if (beforeStart == NULL) {
        return afterStart;
    } else {
        head = beforeStart;
        while (beforeStart->next) {
            beforeStart = beforeStart->next;
        }
        beforeStart->next = afterStart;
        return head;
    }
}