partition linked list

partition_linked_list_2_4.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Question.
# Write code to partition a linked list around a value x, such that nodes less than x come before
# all nodes greater than or equal to x. If x is contained within the list, the values of x only need
# to be after the elements less than x(see below). The partition element x can appear anywhere in the
# "right partition"; it does not need to apeear between the left and right partitions.
# Example
# Input : 3 -> 5 -> 8 -> 5 -> 10 -> 2 -> 1[partition=5]
# Output: 3 -> 1 -> 2 -> 10 -> 5 -> 5 -> 8


class Node:
  def __init__(self, value, next=None):
    self.value = value
    self.next_ptr = next

  def __str__(self):
    return "%d" % (self.value)

class LinkedList:
  def __init__(self):
    self.edge_node = None

  def add_node(self, node):
    node.next_ptr = self.edge_node
    self.edge_node = node

  def add(self, node_value):
    new_node = Node(node_value, self.edge_node)
    self.add_node(new_node)

    return new_node

  def delete_node(self, node):
    pre_node = None
    cur_node = self.edge_node
    while cur_node:
      if cur_node == node:
        if cur_node == self.edge_node:
          self.edge_node = cur_node.next_ptr
        else:
          pre_node.next_ptr = cur_node.next_ptr
      else:
        pre_node = cur_node

      cur_node = cur_node.next_ptr

  def partition(self, x):
    # Input : 3 -> 5 -> 8 -> 5 -> 10 -> 2 -> 1[partition=5]
    # Output: 3 -> 1 -> 2 -> 10 -> 5 -> 5 -> 8
    cur_node = self.edge_node
    while cur_node:
      if cur_node.value < x:
        self.add(cur_node.value)
        self.delete_node(cur_node)
      cur_node = cur_node.next_ptr

  def __str__(self):
    node_value_arr = []
    node = self.edge_node
    while node:
      node_value_arr.append(str(node))
      node = node.next_ptr

    return " -> ".join(node_value_arr)

if __name__ == '__main__':
  # Setup linkedlist
  linked_list = LinkedList()
  node_1 = linked_list.add(1)
  node_2 = linked_list.add(2)
  node_10 = linked_list.add(10)
  node_5 = linked_list.add(5)
  node_8 = linked_list.add(8)
  node_5 = linked_list.add(5)
  node_3 = linked_list.add(3)

  print "Before partition"
  print linked_list

  linked_list.partition(5)

  print "After partition"
  print linked_list