Implemented n sum (keon#365)

Ernest · goswami-rahul · commit c9cfdef3a04e · 2018-06-29T13:30:23.000+05:30
* Implement N sum

* Add N sum test

* Add docstring

* Add n num link

* Rearrange code

* Fix import error

* Move functions to inner

* Fix coding style

* Rename

* Separate logic

* Add advanced usage example

* Add test

* Fix error: cannot hash list object when user's nums is a list of list

* Add parameters docstring
diff --git a/README.md b/README.md
@@ -73,6 +73,7 @@ If you want to uninstall algorithms, it is as simple as:
     - [top_1](algorithms/arrays/top_1.py)
     - [two_sum](algorithms/arrays/two_sum.py)
     - [move_zeros](algorithms/arrays/move_zeros.py)
+    - [n_sum](algorithms/arrays/n_sum.py)
 - [backtrack](algorithms/backtrack)
     - [general_solution.md](algorithms/backtrack/)
     - [anagram](algorithms/backtrack/anagram.py)
diff --git a/algorithms/arrays/__init__.py b/algorithms/arrays/__init__.py
@@ -15,3 +15,4 @@
 from .top_1 import *
 from .two_sum import *
 from .limit import *
+from .n_sum import *
diff --git a/algorithms/arrays/n_sum.py b/algorithms/arrays/n_sum.py
@@ -0,0 +1,129 @@
+"""
+Given an array of n integers, are there elements a, b, .. , n in nums
+such that a + b + .. + n = target?
+
+Find all unique triplets in the array which gives the sum of target.
+
+Example:
+    basic:
+        Given:
+            n = 4, nums = [1, 0, -1, 0, -2, 2], target = 0,
+        return [[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]
+
+    advanced:
+        Given:
+            n = 2
+            nums = [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]]
+            taget = -5
+            def sum(a, b):
+                return [a[0] + b[1], a[1] + b[0]]
+            def compare(num, taget):
+                if num[0] < taget:
+                    return -1
+                elif if num[0] > taget:
+                    return 1
+                else:
+                    return 0
+        return [[-9, 5], [8, 4]]
+        because -9 + 4 = -5
+"""
+
+
+def n_sum(n, nums, target, **kv):
+    """
+    n: int
+    nums: list[object]
+    target: object
+    sum_closure: function, optional
+                 Given two elements of nums, return sum of both.
+    compare_closure: function, optional
+                     Given one object of nums and target, return one of -1, 1, or 0.
+    same_closure: function, optional
+                  Given two object of nums, return bool.
+    return: list[list[object]]
+
+    Note:
+    1. type of sum_closure's return should be same as type of compare_closure's first param
+    """
+
+    def sum_closure_default(a, b):
+        return a + b
+
+    def compare_closure_default(num, taget):
+        if num < taget:
+            return -1
+        elif num > taget:
+            return 1
+        else:
+            return 0
+
+    def same_closure_default(a, b):
+        return a == b
+
+    def n_sum(n, nums, target):
+        if n == 2:
+            results = two_sum(nums, target)
+        else:
+            results = []
+            prev_num = None
+            for index, num in enumerate(nums):
+                if prev_num is not None and \
+                   same_closure(prev_num, num):
+                    continue
+                prev_num = num
+                n_minus1_results = n_sum(n - 1,
+                                         nums[index + 1:],
+                                         target - num)
+                n_minus1_results = append_elem_to_each_list(num,
+                                                            n_minus1_results)
+                results += n_minus1_results
+        return union(results)
+
+    def two_sum(nums, target):
+        nums.sort()
+        lt = 0
+        rt = len(nums) - 1
+        results = []
+        while lt < rt:
+            sum_ = sum_closure(nums[lt], nums[rt])
+            flag = compare_closure(sum_, target)
+            if flag == -1:
+                lt += 1
+            elif flag == 1:
+                rt -= 1
+            else:
+                results.append(sorted([nums[lt], nums[rt]]))
+                lt += 1
+                rt -= 1
+                while (lt < len(nums) and
+                       same_closure(nums[lt - 1], nums[lt])):
+                    lt += 1
+                while (0 <= rt and
+                       same_closure(nums[rt], nums[rt + 1])):
+                    rt -= 1
+        return results
+
+    def append_elem_to_each_list(elem, container):
+        results = []
+        for elems in container:
+            elems.append(elem)
+            results.append(sorted(elems))
+        return results
+
+    def union(duplicate_results):
+        results = []
+
+        if len(duplicate_results) != 0:
+            duplicate_results.sort()
+            results.append(duplicate_results[0])
+            for result in duplicate_results[1:]:
+                if results[-1] != result:
+                    results.append(result)
+
+        return results
+
+    sum_closure = kv.get('sum_closure', sum_closure_default)
+    same_closure = kv.get('same_closure', same_closure_default)
+    compare_closure = kv.get('compare_closure', compare_closure_default)
+    nums.sort()
+    return n_sum(n, nums, target)
diff --git a/tests/test_array.py b/tests/test_array.py
@@ -15,7 +15,8 @@
     max_ones_index,
     trimmean,
     top_1,
-    limit
+    limit,
+    n_sum
 )
 
 import unittest
@@ -309,28 +310,51 @@ def test_two_sum(self):
         self.assertTupleEqual((0, 3), two_sum([-3, 5, 2, 3, 8, -9], target=0))
 
         self.assertIsNone(two_sum([-3, 5, 2, 3, 8, -9], target=6))
-        
+
+
 class TestTrimmean(unittest.TestCase):
 
     def test_trimmean(self):
 
         self.assertEqual(trimmean([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 20), 5.5)
         self.assertEqual(trimmean([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 20), 6.0)
-        
+
+
 class TestTop1(unittest.TestCase):
-    
+
     def test_top_1(self):
         self.assertListEqual(top_1([1 , 1, 2, 2, 3]), [1, 2])
         self.assertListEqual(top_1([1, 2, 3, 324, 234, 23, 23, 1, 23, 23]), [23])
 
+
 class TestLimit(unittest.TestCase):
-    
+
     def test_limit(self):
         self.assertListEqual(limit([1, 2, 3, 4, 5], 2, 4), [2, 3, 4])
         self.assertListEqual(limit([1, 2, 3, 4, 5], 2), [2, 3, 4, 5])
         self.assertListEqual(limit([1, 2, 3, 4, 5], None, 4), [1, 2, 3, 4])
 
-        
+
+class TestNSum(unittest.TestCase):
+
+    def test_n_sum(self):
+        self.assertEqual(n_sum(2, [-3, 5, 2, 3, 8, -9], 6), [])  # noqa: E501
+        self.assertEqual(n_sum(3, [-5, -4, -3, -2, -1, 0, 1, 2, 3], 0), sorted([[-5,2,3],[-2,0,2],[-4,1,3],[-3,1,2],[-1,0,1],[-2,-1,3],[-3,0,3]]))  # noqa: E501
+        self.assertEqual(n_sum(3, [-1,0,1,2,-1,-4], 0), sorted([[-1,-1,2],[-1,0,1]]))  # noqa: E501
+        self.assertEqual(n_sum(4, [1, 0, -1, 0, -2, 2], 0), sorted([[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]))  # noqa: E501
+        self.assertEqual(n_sum(4, [7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 6, 4, -3, -2], 10), sorted([[-6, 2, 7, 7], [-6, 3, 6, 7], [-6, 4, 5, 7], [-6, 4, 6, 6], [-5, 1, 7, 7], [-5, 2, 6, 7], [-5, 3, 5, 7], [-5, 3, 6, 6], [-5, 4, 4, 7], [-5, 4, 5, 6], [-4, 0, 7, 7], [-4, 1, 6, 7], [-4, 2, 5, 7], [-4, 2, 6, 6], [-4, 3, 4, 7], [-4, 3, 5, 6], [-4, 4, 4, 6], [-3, -1, 7, 7], [-3, 0, 6, 7], [-3, 1, 5, 7], [-3, 1, 6, 6], [-3, 2, 4, 7], [-3, 2, 5, 6], [-3, 3, 4, 6], [-3, 4, 4, 5], [-2, -2, 7, 7], [-2, -1, 6, 7], [-2, 0, 5, 7], [-2, 0, 6, 6], [-2, 1, 4, 7], [-2, 1, 5, 6], [-2, 2, 3, 7], [-2, 2, 4, 6], [-2, 3, 4, 5], [-1, 0, 4, 7], [-1, 0, 5, 6], [-1, 1, 3, 7], [-1, 1, 4, 6], [-1, 2, 3, 6], [-1, 2, 4, 5], [-1, 3, 4, 4], [0, 1, 2, 7], [0, 1, 3, 6], [0, 1, 4, 5], [0, 2, 3, 5], [0, 2, 4, 4], [1, 2, 3, 4]]))  # noqa: E501
+
+        self.assertEqual(n_sum(2, [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]], 0,  # noqa: E501
+                               sum_closure=lambda a, b: a[0] + b[0]),  # noqa: E501
+                         [[[-3, 0], [3, 3]], [[-2, 1], [2, 2]]])  # noqa: E501
+        self.assertEqual(n_sum(2, [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]], [0, 3],  # noqa: E501
+                               sum_closure=lambda a, b: [a[0] + b[0], a[1] + b[1]],  # noqa: E501
+                               same_closure=lambda a, b: a[0] == b[0] and a[1] == b[1]),  # noqa: E501
+                         [[[-3, 0], [3, 3]], [[-2, 1], [2, 2]]])  # noqa: E501
+        self.assertEqual(n_sum(2, [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]], -5,  # noqa: E501
+                               sum_closure=lambda a, b: [a[0] + b[1], a[1] + b[0]],  # noqa: E501
+                               compare_closure=lambda a, b: -1 if a[0] < b else 1 if a[0] > b else 0),  # noqa: E501
+                         [[[-9, 5], [8, 4]]])  # noqa: E501
 
 
 if __name__ == '__main__':
