LeetCode OJ--Partition List

#include <iostream>
using namespace std;
 
struct ListNode {
     int val;
     ListNode *next;
     ListNode(int x) : val(x), next(NULL) {}
 };
 
class Solution {
public:
    ListNode *partition(ListNode *head, int x) {
        ListNode *head1 = NULL,*head2 = NULL,*tail1 = NULL,*tail2 = NULL,*p = NULL;
        p = head;
        bool flag1 = 0,flag2 = 0;
        while(p)
        {
            if(p->val<x)
            {
                if(flag1 == 0)
                {
                    head1 = tail1 = p;
                    flag1 = 1;
                }
                else
                {
                    tail1->next = p;
                    tail1 = p;
                }
            }
            else
            {
                if(flag2 == 0)
                {
                    head2 = tail2 = p;
                    flag2 = 1;
                }
                else
                {
                    tail2->next = p;
                    tail2 = p;
                }
            }
            p = p->next;
        }
        if(head1 == tail1 && tail1 == NULL)
            return head2;
        tail1->next = head2;
        if(tail2)
        tail2->next = NULL;
        return head1;

    }
};

int main()
{
    ListNode *n1 = new ListNode(1);
    ListNode *n2 = new ListNode(4);
    ListNode *n3 = new ListNode(3);
    ListNode *n4 = new ListNode(2);
    ListNode *n5 = new ListNode(5);
    ListNode *n6 = new ListNode(2);
    n1->next = n2;
    n2->next = n3;
    n3->next = n4;
    n4->next = n5;
    n5->next = n6;
    ListNode *ans;
    Solution myS;
    ans = myS.partition(NULL,6);
    return 0;
}