3sum problem

3sum.py

'''
Given an array S of n integers, are there elements a, b, c in S such that
a + b + c = 0? Find all unique triplets in the array which gives the sum of
zero. Note, the solution must run in O(n^2logn) or O(n^2) time or better and
can use at most O(n) extra memory.

For example, given array S = [-1, 0, 1, 2, -1, -4],

A solution set is:
[
  [-1, 0, 1],
  [-1, -1, 2]
]


Run this by opening your terminal:
$ python 3sum.py

It was tested on Python 2.7
'''


def main(array):
  # Sort the array
  print (array)
  array = sort(array)
  print (array)

  foundCases = []

  # -2 cause we need three values
  for i in range(0, len(array)):
    # i is set

    # if duplicate i value
    if i > 0 and array[i] == array[i-1]:
      continue


    j = i + 1

    # k is from higher bound
    k = len(array) - 1


    count = 0

    # j and k never cross so n is the for loop
    # and this loop makes it n square
    while(j < k):
      # Find the total of the triplet
      total = array[i] + array[j] + array[k]

      # If total is greater than 0 move the highest value down
      if total > 0:
        k -= 1
        continue

      if total == 0:
        successfulTriplet = [array[i], array[j], array[k]]
        foundCases.append(successfulTriplet)

      j += 1
      while (j < k and array[j] == array[j-1]):
        j += 1

  print (foundCases)





# Taken from: http://stackoverflow.com/questions/18262306/quicksort-with-python
# My implementation of quicksort is in Java here:
# https://github.com/keithweaver/technical-interview-prep/blob/master/sorting.md
def sort(array):
    less = []
    equal = []
    greater = []

    if len(array) > 1:
        pivot = array[0]
        for x in array:
            if x < pivot:
                less.append(x)
            if x == pivot:
                equal.append(x)
            if x > pivot:
                greater.append(x)
        # Don't forget to return something!
        return sort(less)+equal+sort(greater)  # Just use the + operator to join lists
    # Note that you want equal ^^^^^ not pivot
    else:  # You need to hande the part at the end of the recursion - when you only have one element in your array, just return the array.
        return array


main([-1, 0, 1, 2, -1, -4]);