Question
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 set must not contain duplicate triplets.
For example, given array S = [-1, 0, 1, 2, -1, -4],
A solution set is:
[
[-1, 0, 1],
[-1, -1, 2]
]
Solution
The idea is to sort the input array nums and then run through all indices of a possible first element nums[first] of a triplet. For each possible first element, the problem becomes finding two sums in the sorted array from first+1 to nums.length-1.
To find 2 sums in a sorted array, we could use a bi-dementioanl search. Set two indecies start and end for the start and end of the search range and compare sum of nums[start] + nums[end] and 0 - nums[first]. When sum == 0 - nums[first], we know a valid 2 sum is found. When sum < 0 - nums[first], increase sum by moving start forward. When sum > 0 - nums[first], decrease sum by moving end backwards.
When moving first, start and end, to avoid duplicate results, they need to be moved to the next different element in nums.
Code
class Solution {
public List<List<Integer>> threeSum(int[] nums) {
List<List<Integer>> result = new ArrayList<>();
Arrays.sort(nums);
for(int i=0; i<nums.length-2; i++) {
int first = nums[i];
List<List<Integer>> sumResult = twoSum(i+1, nums, -first);
for(List<Integer> res : sumResult) {
res.add(0, first);
}
result.addAll(sumResult);
// To reduce the duplicate, look for the last same element
// next i++ will move i to the next different element
while(nums[i] == nums[i+1] && i<nums.length-2) i++;
}
return result;
}
public List<List<Integer>> twoSum(int start, int[] nums, int target) {
List<List<Integer>> twoSums = new ArrayList<>();
int l = start;
int r = nums.length-1;
while(l<r) {
int sum = nums[l] + nums[r];
if(sum == target) {
List<Integer> pair = new ArrayList<>();
pair.add(nums[l]);
pair.add(nums[r]);
twoSums.add(pair);
# find the next different nums[l]
// look for the next different element
l++;
while(l<r && nums[l] == nums[l-1]) l++;
# find the next different nums[r]
r--;
while(l<r && nums[r] == nums[r+1]) r--;
}
else if(sum > target) {
# find the next different nums[r]
r--;
while(l<r && nums[r] == nums[r+1]) r--;
}
else {
# find the next different nums[l]
l++;
while(l<r && nums[l] == nums[l-1]) l++;
}
}
return twoSums;
}
}
Performance
Sorting the array will use O(logn) time. And to find the three sums, the outer loop for iterating first takes O(n). The inner loop for finding all the two sums takes O(n). So overall it takes O(logn) + O(n) * O(n) = O(n^2).
For space complexity, it depends on how many triples can be found in nums.
K Sum solution
The above solution could be easily extended to finding k sum. The key idea is to use recursion to reduce the k sum problem to k-1 sum. And when it is reduced to 2 sum, use the bi-dimensional search method to get all the two sums.
Code
class Solution {
public List<List<Integer>> fourSum(int[] nums, int target) {
Arrays.sort(nums);
// Use 4sum as example
return kSum(nums, target, 4, 0);
}
public List<List<Integer>> kSum(int[] nums, int target, int k, int index) {
if(k==2) return twoSum(nums, target, index);
List<List<Integer>> results = new ArrayList();
if(nums.length < k) return new ArrayList();
for(int i=index; i<nums.length-k+1; i++) {
List<List<Integer>> kMinusOneSumResults = kSum(nums, target-nums[i], k-1, i+1);
for(List<Integer> res : kMinusOneSumResults) {
res.add(0, nums[i]);
results.add(res);
}
while(i<nums.length-k+1 && nums[i+1] == nums[i]) i++;
}
return results;
}
public List<List<Integer>> twoSum(int[] nums, int target, int index) {
List<List<Integer>> results = new ArrayList();
int l = index, r = nums.length-1;
while(l<r) {
if(nums[l] + nums[r] == target) {
List<Integer> res = new ArrayList();
res.add(nums[l]);
res.add(nums[r]);
results.add(res);
l++;
while(l<r && nums[l] == nums[l-1]) l++;
r--;
while(l<r && nums[r] == nums[r+1]) r--;
}
else if(nums[l] + nums[r] < target) {
l++;
}
else {
r--;
}
}
return results;
}
}
Performance
Similar to 3 sum, the time complexity will be O(n^k-1). And space complexity depends on how many k sums can be found in nums