partition_list.py

"""
Definition of ListNode
class ListNode(object):

def __init__(self, val, next=None):
    self.val = val
    self.next = next
"""
class Solution:
    """
    @param head: The first node of linked list.
    @param x: an integer
    @return: a ListNode
    """
    def partition(self, head, x):
        # write your code here
        if head is None:
            return None

        small_dummy = ListNode(0)
        big_dummy = ListNode(0)

        small_curr = small_dummy
        big_curr = big_dummy

        node = head
        while node is not None:
            if node.val < x:
                small_curr.next = node
                small_curr = node
            else:
                big_curr.next = node
                big_curr = node
            node = node.next

        big_curr.next = None # TODO correct: close the list sometimes!!!
        small_curr.next = big_dummy.next
        return small_dummy.next