暴力枚举——妖梦拼木棒

暴力枚举——妖梦拼木棒

题目背景

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

题目描述

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

答案对 10⁹ + 7 取模

输入格式

第一行一个整数 n。
第二行 n 个整数，第 i 个整数 ai 代表第 i 根木棒的长度。

输出格式

一行一个整数代表答案

输入输出样例

输入样例

4
1 1 2 2

输出样例

1

说明/提示

数据规模与约定

• 对于 30% 的数据，保证 n ≤ 5 × 10³。
• 对于 100% 的数据，保证 1 ≤ n ≤ 10⁵,0 ≤ ai ≤ 5 × 10³。

分析

假设选取的四根木棒的长度分别为 a,b,c,d，
因为要得到正三角形，所以可以得到等式：a = b = c + d (c ≤ d)，
将长度为 i 的木棒根数存放在 s[i] 中，
若 c = d，则这种情况的个数为：

若 c 不等于 d，则这种情况的个数为：

将结果相加取余即可得到最后的答案。

#include<bits/stdc++.h>
using namespace std;
const int mod=1e9+7;
int f[100005],s[5005],n,cnt=0;

int main()
{
	cin>>n;
	memset(s,0,sizeof(s));
	for(int i=0;i<n;i++){
		cin>>f[i];
		s[f[i]]++;	// 记录每种长度木棒的根数 
	}
	for(int a=2;a<=5000;a++){
		for(int c=1;c<=a/2;c++){ // c 最大不会超过 a/2 
			int d=a-c;
			if(c==d&&s[a]>=2&&s[c]>=2){
				cnt+=(s[a]*(s[a]-1)/2)*(s[c]*(s[c]-1)/2)%mod;
			} 
			else if(c!=d&&s[a]>=2&&s[c]>=1&&s[d]>=1){
				cnt+=((s[a]*(s[a]-1))/2)*s[c]*s[d]%mod;
			}
			cnt=cnt%mod;
		}
	}
	cout<<cnt<<endl;
	return 0;
}