洛谷 小 Y 拼木棒 贪心

题目背景

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

题目描述

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

答案对 109+7 取模。

输入格式

第一行一个整数 n。

第二行往下 n 行，每行 1 个整数，第 i 个整数 ai​ 代表第 i 根木棒的长度。

输出格式

一行一个整数代表答案。

输入输出样例

输入 #1复制

4 
1
1
2
2

输出 #1复制

1

说明/提示

数据规模与约定

• 对于 30% 的数据，保证 n≤5×103。
• 对于 100% 的数据，保证 1≤n≤105，1≤ai​≤5×103。

代码：

#include <bits/stdc++.h>
#define MX 100005
using namespace std;
const int mod = 1e9+7;
int main() {
long long int n,cnt = 0,ant = 0;
cin>>n;
int a[MX],b[MX];
int f[MX] = {0};
for(int i = 1;i <= n;i++)
{
cin>>a[i];
f[a[i]]++;
if(f[a[i]] == 1)
{
b[++ant] = a[i];
}
}
sort(b+1,b+ant+1);
for(int i = 1;i <= ant;i++)
{
if(f[b[i]] >= 2)
{
for(int j = 1;j <= i;j++)
{
if(b[j] > b[i] /2)break;
if(f[b[i] - b[j]] >= 1)
{
if(b[i] - b[j] != b[j])
{
cnt =(cnt + (f[b[i]] *(f[b[i]] - 1)/2*f[b[j]]*(f[b[i] - b[j]]))%mod) % mod;
}
else if(f[b[j]]>=2)
{
cnt = (cnt + ((f[b[i]] *(f[b[i]] - 1)/2 * f[b[j]]*(f[b[j]] - 1)/2)%mod))%mod;
}
}
}
}
}
cout<<cnt;
return 0;
}