2.4-Partition.py

# CTCI 2.4
# Partition

import unittest
#from LinkedList import LinkedList

# My Solution
    
def partition(head, pivot):
  a_head, a_tail = None, None
  b_head, b_tail = None, None
  node = head
  while node:
    if node.data < pivot:
      if a_head:
        a_tail.next, a_tail = node, node
      else:
        a_head, a_tail = node, node
    else:
      if b_head:
        b_tail.next, b_tail = node, node
      else:
        b_head, b_tail = node, node
    node = node.next
  a_tail.next = b_head
  return a_head

#-------------------------------------------------------------------------------
# CTCI Solution

def apartition(ll, x):
    current = ll.tail = ll.head

    while current:
        nextNode = current.next
        current.next = None
        if current.value < x:
            current.next = ll.head
            ll.head = current
        else:
            ll.tail.next = current
            ll.tail = current
        current = nextNode
        
    # Error check in case all nodes are less than x
    if ll.tail.next is not None:
        ll.tail.next = None

#-------------------------------------------------------------------------------
#Testing

class Node():
  def __init__(self, data, next=None):
    self.data, self.next = data, next
  
  def __str__(self):
    string = str(self.data)
    if self.next:
      string += ',' + str(self.next)
    return string

class Test(unittest.TestCase):
  def test_partition(self):
    head1 = Node(7,Node(2,Node(9,Node(1,Node(6,Node(3,Node(8)))))))
    head2 = partition(head1, 6)
    self.assertEqual(str(head2), "2,1,3,7,9,6,8")
    head3 = partition(head2, 7)
    self.assertEqual(str(head3), "2,1,3,6,7,9,8")

if __name__ == "__main__":
  unittest.main()