LeetCode。
关于K-sum问题,最低可以降低到O(nk-1)的复杂度,下面这个解法是在Discuss中看到的关于K-sum的一个通用解法:
首先是原题中的4Sum问题:
1 class Solution{ 2 public List<List<Integer>> fourSum(int[] nums, int target) { 3 ArrayList<List<Integer>> res = new ArrayList<List<Integer>>(); 4 int len = nums.length; 5 if (nums == null || len < 4) 6 return res; 7 8 Arrays.sort(nums); 9 10 int max = nums[len - 1]; 11 if (4 * nums[0] > target || 4 * max < target) 12 return res; 13 14 int i, z; 15 for (i = 0; i < len; i++) { 16 z = nums[i]; 17 if (i > 0 && z == nums[i - 1])// avoid duplicate 18 continue; 19 if (z + 3 * max < target) // z is too small 20 continue; 21 if (4 * z > target) // z is too large 22 break; 23 if (4 * z == target) { // z is the boundary 24 if (i + 3 < len && nums[i + 3] == z) 25 res.add(Arrays.asList(z, z, z, z)); 26 break; 27 } 28 29 threeSumForFourSum(nums, target - z, i + 1, len - 1, res, z); 30 } 31 32 return res; 33 } 34 35 /* 36 * Find all possible distinguished three numbers adding up to the target 37 * in sorted array nums[] between indices low and high. If there are, 38 * add all of them into the ArrayList fourSumList, using 39 * fourSumList.add(Arrays.asList(z1, the three numbers)) 40 */ 41 public void threeSumForFourSum(int[] nums, int target, int low, int high, ArrayList<List<Integer>> fourSumList, 42 int z1) { 43 if (low + 1 >= high) 44 return; 45 46 int max = nums[high]; 47 if (3 * nums[low] > target || 3 * max < target) 48 return; 49 50 int i, z; 51 for (i = low; i < high - 1; i++) { 52 z = nums[i]; 53 if (i > low && z == nums[i - 1]) // avoid duplicate 54 continue; 55 if (z + 2 * max < target) // z is too small 56 continue; 57 58 if (3 * z > target) // z is too large 59 break; 60 61 if (3 * z == target) { // z is the boundary 62 if (i + 1 < high && nums[i + 2] == z) 63 fourSumList.add(Arrays.asList(z1, z, z, z)); 64 break; 65 } 66 67 twoSumForFourSum(nums, target - z, i + 1, high, fourSumList, z1, z); 68 } 69 70 } 71 72 /* 73 * Find all possible distinguished two numbers adding up to the target 74 * in sorted array nums[] between indices low and high. If there are, 75 * add all of them into the ArrayList fourSumList, using 76 * fourSumList.add(Arrays.asList(z1, z2, the two numbers)) 77 */ 78 public void twoSumForFourSum(int[] nums, int target, int low, int high, ArrayList<List<Integer>> fourSumList, 79 int z1, int z2) { 80 81 if (low >= high) 82 return; 83 84 if (2 * nums[low] > target || 2 * nums[high] < target) 85 return; 86 87 int i = low, j = high, sum, x; 88 while (i < j) { 89 sum = nums[i] + nums[j]; 90 if (sum == target) { 91 fourSumList.add(Arrays.asList(z1, z2, nums[i], nums[j])); 92 93 x = nums[i]; 94 while (++i < j && x == nums[i]) // avoid duplicate 95 ; 96 x = nums[j]; 97 while (i < --j && x == nums[j]) // avoid duplicate 98 ; 99 } 100 if (sum < target) 101 i++; 102 if (sum > target) 103 j--; 104 } 105 return; 106 } 107 }
然后推广到K-sum问题后,还用到了Backtracking,可以说很elegant了:
1 List<List<Integer>> kSum_Trim(int[] a, int target, int k) { 2 List<List<Integer>> result = new ArrayList<>(); 3 if (a == null || a.length < k || k < 2) return result; 4 Arrays.sort(a); 5 kSum_Trim(a, target, k, 0, result, new ArrayList<>()); 6 return result; 7 } 8 9 void kSum_Trim(int[] a, int target, int k, int start, List<List<Integer>> result, List<Integer> path) { 10 int max = a[a.length - 1]; 11 if (a[start] * k > target || max * k < target) return; 12 13 if (k == 2) { // 2 Sum 14 int left = start; 15 int right = a.length - 1; 16 while (left < right) { 17 if (a[left] + a[right] < target) left++; 18 else if (a[left] + a[right] > target) right--; 19 else { 20 result.add(new ArrayList<>(path)); 21 result.get(result.size() - 1).addAll(Arrays.asList(a[left], a[right])); 22 left++; right--; 23 while (left < right && a[left] == a[left - 1]) left++; 24 while (left < right && a[right] == a[right + 1]) right--; 25 } 26 } 27 } 28 else { // k Sum 29 for (int i = start; i < a.length - k + 1; i++) { 30 if (i > start && a[i] == a[i - 1]) continue; 31 if (a[i] + max * (k - 1) < target) continue; 32 if (a[i] * k > target) break; 33 if (a[i] * k == target) { 34 if (a[i + k - 1] == a[i]) { 35 result.add(new ArrayList<>(path)); 36 List<Integer> temp = new ArrayList<>(); 37 for (int x = 0; x < k; x++) temp.add(a[i]); 38 result.get(result.size() - 1).addAll(temp); // Add result immediately. 39 } 40 break; 41 } 42 path.add(a[i]); 43 kSum_Trim(a, target - a[i], k - 1, i + 1, result, path); 44 path.remove(path.size() - 1); // Backtracking 45 } 46 } 47 }