3Sum Problem

3sum.py

"""
    Problem Statement
    Given a set of numbers M and another number k
    Find if there exists 3 numbers a,b,c in M such that
    a+b+c = k
    in complexity O(N^2)
    Solution
    This code solves the problem in O(N^2). However, following assumptions had been made
    Set comprises of positive integers.
    Note - the code can be customized to handle negative numbers and decimals
    Bonus Tip - When wiki says that the best case complexity is O(N^2), then its most probably correct
    Refer - http://en.wikipedia.org/wiki/3SUM
"""

import sys


def merge(left, right):
    result = []
    i ,j = 0, 0
    while i < len(left) and j < len(right):
        if left[i] <= right[j]:
            result.append(left[i])
            i += 1
        else:
            result.append(right[j])
            j += 1
    result += left[i:]
    result += right[j:]
    return result

def mergesort(list):
    """
    mergesort function
    Complexity - O(NlogN)
    Ref - http://en.literateprograms.org/Merge_sort_(Python)
    """
    if len(list) < 2:
        return list
    else:
        middle = len(list) / 2
        left = mergesort(list[:middle])
        right = mergesort(list[middle:])
        return merge(left, right)

def binary_search(a, x, lo=0, hi=None):
    """
    Complexity - O(NlogN)
    """
    if hi is None:
        hi = len(a)
    while lo < hi:
        mid = (lo+hi)//2
        midval = a[mid]
        if midval < x:
            lo = mid+1
        elif midval > x: 
            hi = mid
        else:
            return mid
    return -1    

def check_input_set(number_set, k):
    """
    checks if in a number list three numbers a,b,c exist such that a+b+c = k
    returns [a,b,c] if exists else None
    Complexity - O(N^2)
    """
    # Sort the number set
    sorted_number_list = mergesort(number_set)
    print 'Sorted Numbers %s' %(sorted_number_list)    
    i = 0
    j = 0
    # Find index of the largest number in sorted list that is < k
    for l, e in reversed(list(enumerate(sorted_number_list))):
        if e < k:
            j = l
            break
    if (j == 0):
        print 'No such number set exists'
        return None
    for i in range(0,j+1):
        a = k - sorted_number_list[i]
        l = i
        m = j
        if a > 0:
            while l < m:
                two_sum = sorted_number_list[l]+sorted_number_list[m]
                if two_sum < a:
                    l = l+1
                elif two_sum > a:
                    m = m-1
                else:
                    b = sorted_number_list[l]
                    c = sorted_number_list[m]
                    return [k-b-c,b,c]
    return None                  
        

def main():
    # Take input of list of numbers from console
    input_flag = True
    k = 0
    number_set = []
    while input_flag:
        number = raw_input('Please enter a positive integer - ')
        try:
            number = int(number)
            number_set.append(number)
        except:
            print 'Please enter a positive integer'
            print sys.exc_info()[0]
        cont_flag = raw_input('Continue y or n - ')
        cont_flag = cont_flag.lower()
        if cont_flag != 'y':
            input_flag = False
    print 'You have entered the following numbers %s' %(number_set)
    k = raw_input('Please enter the value of k - ')
    try:
        k = int(k)
    except:
        print 'Please enter the value'
        print sys.exc_info()[0]
    result_set = check_input_set(number_set, k)
    print result_set
    
if __name__ == "__main__":
    main()