【nowcoder】链表分割

牛客题目链接 链表分割

/*
struct ListNode {
    int val;
    struct ListNode *next;
    ListNode(int x) : val(x), next(NULL) {}
};*/
#include <cstddef>
class Partition {
  public:
    ListNode* partition(ListNode* pHead, int x) {
        if (pHead != NULL) {
            // 1.两个链表，一个放<x值的节点，一个放>=x值的节点
            ListNode* head1, *head2, *tail1, *tail2;
            head1 = head2 = tail1 = tail2 = NULL; 
            ListNode* cur = pHead;
            while (cur != NULL) {
                if (cur->val < x) {
                    if (head1 == NULL) {
                        head1 = tail1 = cur;
                    } else {
                        tail1->next = cur;
                        tail1 = tail1->next;
                    }
                } else {
                    if (head2 == NULL) {
                        head2 = tail2 = cur;
                    } else {
                        tail2->next = cur;
                        tail2 = tail2->next;
                    }
                }
                cur = cur->next;
            }

            // 2.把两个链表连起来，不过有几种情况：
            if (head1 != NULL && head2 != NULL) {
                tail1->next = head2;
                tail2->next = NULL;
                return head1;
            } else if (head1 != NULL && head2 == NULL) {
                return head1;
            } else if (head1 == NULL && head2 != NULL) {
                return head2;
            } // 不用考虑head1和head2同时为NULL的情况，这时pHead就是NULL，代码执行进不来的
        }
        return NULL;
    }
};