[leetcode] Partition List

#include "iostream"
using namespace std;

/**Definition for singly-linked list.*/

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

class Solution {
public:
ListNode *partition(ListNode *head, int x) {

if (head==NULL) {
return NULL;
}

//cout<<head->val<<endl;

ListNode *small=new ListNode(0);
ListNode *big=new ListNode(0);
ListNode *p=small;
ListNode *q=big;
ListNode *current=head;

//cout<<current->val<<endl;

while (current!=NULL) {
if (current->val<x) {
p->next=current;
current=current->next;//切记，不能使用current++！！！
p=p->next;
}else{
q->next=current;
current=current->next;
q=q->next;
}
}

if (p==small||q==big) {//都比x大or小
return head;
}

//连接两个串

p->next=big->next;
q->next=NULL;
current=small->next;

//记得释放内存
free(small);
free(big);

return current;

}
};

int main() {
ListNode *t=new ListNode(1);
t->next=new ListNode(2);
cout<<t->val<<endl;
Solution so;
ListNode *res= so.partition(t, 2);
cout<<res->val<<res->next->val<<endl;
return 0;
}