算法学习之一：3SUM变体

算法题目

在一个包含n个整数的数组S中，是否存在三个元素(a,b,c)，使得a+b=c？如果存在，给出所有满足a+b=c的(a,b,c)集合，不能重复。

注意：

1. 三元素组合(a,b,c)，要求(a,b)是升序的，即a<=b
2. 结果集不能重复

例子：

数组S = {1, 1, 2, -1, -4}，

结果集为：
{-1, 2, 1}
{1, 1, 2}

思路

这是我面试一个动漫公司遇到的题，是经典3SUM问题的一种变体，整体思路跟3SUM一样，但略作改变。

第一步，排序，排成降序的，如果S = {1, 1, 2, -1, -4}，则排序后为 {2, 1, 1, -1, -4}，这一步的时间复杂度为O(nlogn)；

第二步，遍历排序后的数组，并取遍历得到的值为c。确定了c以后，接下来问题就变成了，在排序后的数组中寻找两个元素之和等于某个数，这是经典的2SUM问题，时间复杂度为O(n)。总共有n次遍历，每一次的时间复杂度为O(n)，因此这一步的时间复杂度为 O(n^2)。

两步总的时间复杂度为O(nlogn) + O(n^2) = O(n^2)。

亮点

我在学习过程中，特别惊艳于，能在O(n)的时间复杂度内，从有序数组中寻找两个元素之和等于某个数。

具体的思路如下：

还是十分巧妙的。

代码

#include <iostream>
#include <vector>
#include <set>
#include <algorithm>
using namespace std;

bool AscendingCompare(int i,int j) { 
    return i > j;
}

class Solution {
public:
	vector<vector<int>> FindSumEqual(vector<int>& array) {
		sort(array.begin(), array.end(), AscendingCompare);
    
		set<vector<int>> results;
		for (int i = 0; i < array.size(); i++) {
			int j = 0;
			int k = array.size() - 1;
			while (j < k) {
				if (array[j] + array[k] == array[i]) {
					results.insert({array[k], array[j], array[i]});
					j++;
					k--;
					while (j < k && array[j] == array[j + 1]) {
						j++;
					}
					while (j < k &&  array[k] == array[k - 1]) {
						k--;
					}
				} else if (array[j] + array[k] > array[i]) {
					j++;
				} else {
					k--;
				}
			}
		}
		vector<vector<int>> results_vector(results.begin(), results.end());

		return results_vector;
	}
}