P3799 妖梦拼木棒

妖梦拼木棒

题目背景

上道题中，妖梦斩了一地的木棒，现在她想要将木棒拼起来。

题目描述

有 \(n\) 根木棒，现在从中选 \(4\) 根，想要组成一个正三角形，问有几种选法？

答案对 \(10^9+7\) 取模。

输入格式

第一行一个整数 \(n\)。

第二行往下 \(n\) 行，每行 \(1\) 个整数，第 \(i\) 个整数 \(a_i\) 代表第 \(i\) 根木棒的长度。

输出格式

一行一个整数代表答案。

样例 #1

样例输入 #1

4 
1
1
2
2

样例输出 #1

1

提示

数据规模与约定

• 对于 \(30\%\) 的数据，保证 \(n \le 5 \times 10^3\)。
• 对于 \(100\%\) 的数据，保证 \(1 \leq n \le 10^5\)，\(1 \le a_i \le 5 \times 10^3\)。

2.题解

2.1 子集枚举

错误思路

看到从n个里面选4个，自然而然想到了子集枚举，但是自己错误估计了范围
由于子集枚举中会用到 1 << n, 而 \(1≤ n ≤10^5\) 所以会产生一个非常大的数，根本无法存储，导致溢出

错误代码

#include<bits/stdc++.h>
using namespace std;

int main(){
	int n;
	long long ans = 0; 
	cin >> n;
	if(n < 4) cout << 0;
	
	vector<int> arr(n);
	for(int i = 0; i < n; i++)
		cin >> arr[i];
	
	long long U = 1 << n;
	for(long long S = 0; S < U; S++){
		if(__builtin_popcount(S) == 4){
			vector<int> tmp;
			for(int i = 0; i < n; i++){
				if(S & (1 << i)) tmp.push_back(arr[i]);
			}
			sort(tmp.begin(), tmp.end());
			if(tmp[0] + tmp[1] == tmp[2] && tmp[2] == tmp[3]) 
				ans++;
//			if((tmp[0] == tmp[1] && tmp[1] == tmp[2]) || (tmp[0] == tmp[1] && tmp[1] == tmp[3]) 
//			|| (tmp[0] == tmp[2] && tmp[2] == tmp[3]) || (tmp[1] == tmp[2] && tmp[2] == tmp[3])	) 
//				ans++;
		}
	} 
//	cout << ans % static_cast<int>(1e9 + 7); // 使用 1e9 表示科学计数法中的 $10^9$，但它是一个 double 类型的浮点数，而不是整数,% 操作符使用时，左右操作数都必须是整数
	cout << ans % 1000000007;	
} 

2.2 循环枚举

思路

由于 0 <= n <= 5000, 且\(1 ≤ a_i ​≤5×10^3\), 火柴长度的范围不是很大，只有5000，所以我们直接用数组存储每种数组出现的次数。
同时记录最小数和最大数，方便后面循环使用。
这里我们最开始的思路是四重循环，分别表示四根火柴棒的长度，但是\(O(n^4)\)时间复杂度过高
我们寻找它们之间的关系，要构成一个等边三角形，首先有两个数一定是相等的，另外两个数的和也和前者相等。
所以我们其实只需要两重循环，分别表示两个作为和的数，并用他们的和去我们的数组中寻找有没有两个以上该长度的火柴棒即可
这里注意一下：
1.对于两个作为和的数，相等和不相等的时候可以分出来讨论一下，也可以不用。
2.第二层循环必定是[i, maxNum],就比如像之前已经找过长度为1和长度为2的火柴棒了，再去找长度为2的火柴棒和长度为1的火柴棒就重复了！！！

代码

#include<bits/stdc++.h>
using namespace std;
const int MOD = 1e9 + 7;

int C(int n, int k){
	if(k == 1) return n;
	else if(k == 2) return n * (n - 1) / 2;
}
int main(){
	int n;
	long long ans = 0; 
	cin >> n;
	if(n < 4) cout << 0;
	
	vector<int> arr(10001);
	int maxNum = 5000, minNum = 0;
	for(int i = 0; i < n; i++){
		int tmp;
		cin >> tmp;
		arr[tmp]++;
		maxNum = max(maxNum, tmp);
		minNum = min(minNum, tmp);
	}
	
	for(int i = minNum; i <= maxNum; i++){
		for(int j = i; j <= maxNum; j++){
			if(i + j > maxNum) break;
			if(i == j){
				if(arr[i << 1] >= 2 && arr[i] >= 2){
					ans = (ans + C(arr[i], 2) * C(arr[i << 1], 2)) % MOD; 					
				}
			}
			else{
				if(arr[i + j] >= 2 && arr[i] >= 1 && arr[j] >= 1){
					ans = (ans + C(arr[i], 1) * C(arr[j], 1) * C(arr[i + j], 2)) % MOD;
				}
			}
		}
	} 
	cout << ans % MOD;	
}