机试题 三数之和变形-三数和赛高数组01 任意

数组a 对任意的i、j、k都能找到l，使得ai + aj + ak = al，那么这个数组就是被称为三数和赛高数组，特别的i、j、k需要满足（1 <= i < j < k <= n, 1 <= l <= n)。
输入：
首先 t (1 <= t <= 1000) 代表t组数据。
然后 每组数据先输入一个n (3 <= n <= 200000), 代表数组的长度为n，接下来是n个数字a1, a2, ... , an (-10^9 <= ai <= 10^9)
输出：
每组数据，如果是三数和赛高数组就输出"YES"，否则输出"NO"

test1
input:

4
3
1 0 -1
5
1 -2 -2 1 -3
6
0 0 0 0 0 0
4
-1 2 -3 4

output:

YES
NO
YES
NO

my code

#include <iostream>
#include <unordered_set>

using namespace std;

void solution() {
	int n;
	cin >> n;
	int num, cout_zero = 0, sum = 0;
	unordered_set<int> negative, positive;
	for (int i = 0; i < n; ++i){
		cin >> num;
		if (num < 0){
			negative.insert(num);
			sum += num;
		} else if (num > 0){
			positive.insert(num);
			sum += num;
		} else cout_zero++;
	}
	bool flag = false;
	if (n == 3) {
		if (sum == 0) flag = true;
		else if (sum < 0) flag = negative.find(sum) != negative.end();
		else flag = positive.find(sum) != positive.end();
	} else if (n == 4 && sum == 0) {
		if (cout_zero == 4 || (positive.size() == 1 && negative.size() == 0)) flag = true;
	} else if (sum == 0) {
		if (cout_zero == n || (cout_zero == n - 2 && sum == 0)) flag = true;
	}
	cout << (flag ? "YES" : "NO") << endl;
	return;
}

int main() {
	int t;
	cin >> t;
	while(t--) solution();
	return 0;
}

analyse