-
Notifications
You must be signed in to change notification settings - Fork 4.7k
Expand file tree
/
Copy pathn_sum.py
More file actions
127 lines (108 loc) · 3.86 KB
/
Copy pathn_sum.py
File metadata and controls
127 lines (108 loc) · 3.86 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
"""
N-Sum
Given an array of integers, find all unique n-tuples that sum to a target
value. Supports custom sum, comparison, and equality closures for advanced
use cases with non-integer elements.
Reference: https://leetcode.com/problems/4sum/
Complexity:
Time: O(n^(k-1)) where k is the tuple size
Space: O(n^(k-1)) for storing results
"""
from __future__ import annotations
from collections.abc import Callable
from typing import Any
def n_sum(
n: int,
nums: list[Any],
target: Any,
**kv: Callable[..., Any],
) -> list[list[Any]]:
"""Find all unique n-tuples in nums that sum to target.
Args:
n: Size of each tuple to find.
nums: List of elements to search.
target: Desired sum for each n-tuple.
**kv: Optional closures:
sum_closure(a, b) - returns sum of two elements.
compare_closure(num, target) - returns -1, 0, or 1.
same_closure(a, b) - returns True if elements are equal.
Returns:
Sorted list of unique n-tuples (as lists) that sum to target.
Examples:
>>> n_sum(3, [-1, 0, 1, 2, -1, -4], 0)
[[-1, -1, 2], [-1, 0, 1]]
"""
def _sum_closure_default(a: Any, b: Any) -> Any:
return a + b
def _compare_closure_default(num: Any, target: Any) -> int:
if num < target:
return -1
elif num > target:
return 1
else:
return 0
def _same_closure_default(a: Any, b: Any) -> bool:
return a == b
def _n_sum_inner(n: int, nums: list[Any], target: Any) -> list[list[Any]]:
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_inner(
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: list[Any], target: Any) -> list[list[Any]]:
nums.sort()
left = 0
right = len(nums) - 1
results = []
while left < right:
current_sum = sum_closure(nums[left], nums[right])
flag = compare_closure(current_sum, target)
if flag == -1:
left += 1
elif flag == 1:
right -= 1
else:
results.append(sorted([nums[left], nums[right]]))
left += 1
right -= 1
while left < len(nums) and same_closure(nums[left - 1], nums[left]):
left += 1
while right >= 0 and same_closure(nums[right], nums[right + 1]):
right -= 1
return results
def _append_elem_to_each_list(
elem: Any, container: list[list[Any]]
) -> list[list[Any]]:
results = []
for elems in container:
elems.append(elem)
results.append(sorted(elems))
return results
def _union(
duplicate_results: list[list[Any]],
) -> list[list[Any]]:
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_inner(n, nums, target)