程序设计思维（CSP模拟题）A - HRZ 的序列

题目

相较于咕咕东，瑞神是个起早贪黑的好孩子，今天早上瑞神起得很早，刷B站时看到了一个序列a，他对这个序列产生了浓厚的兴趣。
他好奇是否存在一个数K，使得一些数加上K，一些数减去K，一些数不变，使得整个序列中所有的数相等。
其中对于序列中的每个位置上的数字，至多只能执行一次加运算或减运算或是对该位置不进行任何操作。
由于瑞神只会刷B站，所以他把这个问题交给了你！

Input
输入第一行是一个正整数t表示数据组数。
接下来对于每组数据，输入的第一个正整数n表示序列a的长度，随后一行有n个整数，表示序列a。

Output
输出共包含t行，每组数据输出一行。对于每组数据，如果存在这样的K，输出"YES"，否则输出“NO”。（输出不包含引号）

Sample input
2
5
1 2 3 4 5
5
1 2 3 4 5

Sample output
NO
NO

思路

对于输入的序列，可能出现如下5种情况：
① 序列出现1种数字：显然存在K。
② 序列出现2种数字：存在K。例如：
输入：1 1 1 2 2
则K=0.5，前面3个1加上0.5，后面2个2减去0.5，得到 1.5 1.5 1.5 1.5 1.5。
③ 序列出现3种数字n₁、n₂、n₃（n₁ < n₂ < n₃），且n₁ + n₃ = 2 * n₂：存在K。例如：
输入：1 1 1 2 2 3 3 3 3
则K=1，前面3个1加上1，后面4个3减去1，得到 2 2 2 2 2 2 2 2 2。
④ 序列出现3种数字n₁、n₂、n₃（n₁ < n₂ < n₃），且n₁ + n₃ != 2 * n₂：不存在K。
⑤ 序列出现3种以上数字：不存在K。

代码

#include <iostream>
#include <stdio.h>
#include <algorithm>

using namespace std;

long long num[4]; // num[0]不用
bool use[4]; // use[0]不用
bool overflow;

int t;
int n;
int cnt;

void reset()
{
	for (int i = 0; i < 4; i++) use[i] = false;
	cnt = 1;
	overflow = false;
}

void input(long long in)
{
	for (int i = 1; i < 4; i++)
		if (use[i] && num[i] == in)return;

	if (cnt < 4) {
		use[cnt] = true;
		num[cnt] = in;
		cnt++;
	}
	else
		overflow = true;

	return;
}

void print()
{
	for (int i = 1; i < 4; i++) {
		cout << "num " << i << ":";
		if (use[i]) cout << num[i] << endl;
		else cout << endl;
	}

	if (overflow) cout << "OVERFLOW!" << endl;
	else cout << "Not overflow" << endl;
}

void printArray()
{
	for (int i = 1; i < 4; i++)
		cout << num[i] << " ";
	cout << endl;
}


int main()
{
	cin >> t;

	for (int i = 0; i < t; i++) {
		cin >> n;
		reset();
		for (int j = 0; j < n; j++) {
			long long in;
			cin >> in;

			input(in);
		}

		//print();

		//printArray();

		sort(num + 1, num + cnt);

		//printArray();

		int min = 1;
		int max = cnt - 1;

		//cout << max << endl;

		bool solve = false;

		if (overflow) solve = false;
		else if (max < 3) solve = true;
		else if (max == 3 && num[1] + num[3] == 2 * num[2]) solve = true;
		else solve = false;

		if(solve) cout << "YES" << endl;
		else cout << "NO" << endl;
	}

	return 0;
}