LeetCode/cpp/partition-list.cpp

/**
 * Definition for singly-linked list.
 * struct ListNode {
 *     int val;
 *     ListNode *next;
 *     ListNode() : val(0), next(nullptr) {}
 *     ListNode(int x) : val(x), next(nullptr) {}
 *     ListNode(int x, ListNode *next) : val(x), next(next) {}
 * };
 */
class Solution {
    ListNode *update_tail(ListNode *tail, ListNode *node) const {
        if (tail != nullptr) {
            tail->next = node;
        }
        return node;
    }

  public:
    ListNode *partition(ListNode *head, int x) {
        ListNode *left_head = nullptr;
        ListNode *right_head = nullptr;

        ListNode *left_tail = nullptr;
        ListNode *right_tail = nullptr;

        while (head != nullptr) {
            auto next_head = head->next;

            if (head->val < x) {
                left_tail = update_tail(left_tail, head);
            } else {
                right_tail = update_tail(right_tail, head);
            }

            if (left_head == nullptr) {
                left_head = left_tail;
            }
            if (right_head == nullptr) {
                right_head = right_tail;
            }

            head = next_head;
        }

        left_tail = update_tail(left_tail, right_head);
        right_tail = update_tail(right_tail, nullptr);

        if (left_head == nullptr) {
            return left_tail;
        }
        return left_head;
    }
};