Skip to content

Commit c9cfdef

Browse files
Ernestgoswami-rahul
authored andcommitted
Implemented n sum (keon#365)
* 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
1 parent a289b6b commit c9cfdef

File tree

4 files changed

+161
-6
lines changed

4 files changed

+161
-6
lines changed

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ If you want to uninstall algorithms, it is as simple as:
7373
- [top_1](algorithms/arrays/top_1.py)
7474
- [two_sum](algorithms/arrays/two_sum.py)
7575
- [move_zeros](algorithms/arrays/move_zeros.py)
76+
- [n_sum](algorithms/arrays/n_sum.py)
7677
- [backtrack](algorithms/backtrack)
7778
- [general_solution.md](algorithms/backtrack/)
7879
- [anagram](algorithms/backtrack/anagram.py)

algorithms/arrays/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@
1515
from .top_1 import *
1616
from .two_sum import *
1717
from .limit import *
18+
from .n_sum import *

algorithms/arrays/n_sum.py

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
"""
2+
Given an array of n integers, are there elements a, b, .. , n in nums
3+
such that a + b + .. + n = target?
4+
5+
Find all unique triplets in the array which gives the sum of target.
6+
7+
Example:
8+
basic:
9+
Given:
10+
n = 4, nums = [1, 0, -1, 0, -2, 2], target = 0,
11+
return [[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]
12+
13+
advanced:
14+
Given:
15+
n = 2
16+
nums = [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]]
17+
taget = -5
18+
def sum(a, b):
19+
return [a[0] + b[1], a[1] + b[0]]
20+
def compare(num, taget):
21+
if num[0] < taget:
22+
return -1
23+
elif if num[0] > taget:
24+
return 1
25+
else:
26+
return 0
27+
return [[-9, 5], [8, 4]]
28+
because -9 + 4 = -5
29+
"""
30+
31+
32+
def n_sum(n, nums, target, **kv):
33+
"""
34+
n: int
35+
nums: list[object]
36+
target: object
37+
sum_closure: function, optional
38+
Given two elements of nums, return sum of both.
39+
compare_closure: function, optional
40+
Given one object of nums and target, return one of -1, 1, or 0.
41+
same_closure: function, optional
42+
Given two object of nums, return bool.
43+
return: list[list[object]]
44+
45+
Note:
46+
1. type of sum_closure's return should be same as type of compare_closure's first param
47+
"""
48+
49+
def sum_closure_default(a, b):
50+
return a + b
51+
52+
def compare_closure_default(num, taget):
53+
if num < taget:
54+
return -1
55+
elif num > taget:
56+
return 1
57+
else:
58+
return 0
59+
60+
def same_closure_default(a, b):
61+
return a == b
62+
63+
def n_sum(n, nums, target):
64+
if n == 2:
65+
results = two_sum(nums, target)
66+
else:
67+
results = []
68+
prev_num = None
69+
for index, num in enumerate(nums):
70+
if prev_num is not None and \
71+
same_closure(prev_num, num):
72+
continue
73+
prev_num = num
74+
n_minus1_results = n_sum(n - 1,
75+
nums[index + 1:],
76+
target - num)
77+
n_minus1_results = append_elem_to_each_list(num,
78+
n_minus1_results)
79+
results += n_minus1_results
80+
return union(results)
81+
82+
def two_sum(nums, target):
83+
nums.sort()
84+
lt = 0
85+
rt = len(nums) - 1
86+
results = []
87+
while lt < rt:
88+
sum_ = sum_closure(nums[lt], nums[rt])
89+
flag = compare_closure(sum_, target)
90+
if flag == -1:
91+
lt += 1
92+
elif flag == 1:
93+
rt -= 1
94+
else:
95+
results.append(sorted([nums[lt], nums[rt]]))
96+
lt += 1
97+
rt -= 1
98+
while (lt < len(nums) and
99+
same_closure(nums[lt - 1], nums[lt])):
100+
lt += 1
101+
while (0 <= rt and
102+
same_closure(nums[rt], nums[rt + 1])):
103+
rt -= 1
104+
return results
105+
106+
def append_elem_to_each_list(elem, container):
107+
results = []
108+
for elems in container:
109+
elems.append(elem)
110+
results.append(sorted(elems))
111+
return results
112+
113+
def union(duplicate_results):
114+
results = []
115+
116+
if len(duplicate_results) != 0:
117+
duplicate_results.sort()
118+
results.append(duplicate_results[0])
119+
for result in duplicate_results[1:]:
120+
if results[-1] != result:
121+
results.append(result)
122+
123+
return results
124+
125+
sum_closure = kv.get('sum_closure', sum_closure_default)
126+
same_closure = kv.get('same_closure', same_closure_default)
127+
compare_closure = kv.get('compare_closure', compare_closure_default)
128+
nums.sort()
129+
return n_sum(n, nums, target)

tests/test_array.py

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@
1515
max_ones_index,
1616
trimmean,
1717
top_1,
18-
limit
18+
limit,
19+
n_sum
1920
)
2021

2122
import unittest
@@ -309,28 +310,51 @@ def test_two_sum(self):
309310
self.assertTupleEqual((0, 3), two_sum([-3, 5, 2, 3, 8, -9], target=0))
310311

311312
self.assertIsNone(two_sum([-3, 5, 2, 3, 8, -9], target=6))
312-
313+
314+
313315
class TestTrimmean(unittest.TestCase):
314316

315317
def test_trimmean(self):
316318

317319
self.assertEqual(trimmean([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 20), 5.5)
318320
self.assertEqual(trimmean([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], 20), 6.0)
319-
321+
322+
320323
class TestTop1(unittest.TestCase):
321-
324+
322325
def test_top_1(self):
323326
self.assertListEqual(top_1([1 , 1, 2, 2, 3]), [1, 2])
324327
self.assertListEqual(top_1([1, 2, 3, 324, 234, 23, 23, 1, 23, 23]), [23])
325328

329+
326330
class TestLimit(unittest.TestCase):
327-
331+
328332
def test_limit(self):
329333
self.assertListEqual(limit([1, 2, 3, 4, 5], 2, 4), [2, 3, 4])
330334
self.assertListEqual(limit([1, 2, 3, 4, 5], 2), [2, 3, 4, 5])
331335
self.assertListEqual(limit([1, 2, 3, 4, 5], None, 4), [1, 2, 3, 4])
332336

333-
337+
338+
class TestNSum(unittest.TestCase):
339+
340+
def test_n_sum(self):
341+
self.assertEqual(n_sum(2, [-3, 5, 2, 3, 8, -9], 6), []) # noqa: E501
342+
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
343+
self.assertEqual(n_sum(3, [-1,0,1,2,-1,-4], 0), sorted([[-1,-1,2],[-1,0,1]])) # noqa: E501
344+
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
345+
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
346+
347+
self.assertEqual(n_sum(2, [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]], 0, # noqa: E501
348+
sum_closure=lambda a, b: a[0] + b[0]), # noqa: E501
349+
[[[-3, 0], [3, 3]], [[-2, 1], [2, 2]]]) # noqa: E501
350+
self.assertEqual(n_sum(2, [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]], [0, 3], # noqa: E501
351+
sum_closure=lambda a, b: [a[0] + b[0], a[1] + b[1]], # noqa: E501
352+
same_closure=lambda a, b: a[0] == b[0] and a[1] == b[1]), # noqa: E501
353+
[[[-3, 0], [3, 3]], [[-2, 1], [2, 2]]]) # noqa: E501
354+
self.assertEqual(n_sum(2, [[-3, 0], [-2, 1], [2, 2], [3, 3], [8, 4], [-9, 5]], -5, # noqa: E501
355+
sum_closure=lambda a, b: [a[0] + b[1], a[1] + b[0]], # noqa: E501
356+
compare_closure=lambda a, b: -1 if a[0] < b else 1 if a[0] > b else 0), # noqa: E501
357+
[[[-9, 5], [8, 4]]]) # noqa: E501
334358

335359

336360
if __name__ == '__main__':

0 commit comments

Comments
 (0)