妖梦拼木棒

题目背景

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

题目描述

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

输入格式

第一行一个整数n

第二行n个整数，a1,a2,……an(0<ai<=5000)，代表每根木棒的长度。

输出格式

一行一个整数，对1e9+7取模

输入输出样例

输入 #1

4 1 1 2 2

输出 #1

1

说明/提示

对于30%的数据 N<=5000

对于100%的数据 N<=100000

by-szc

 

【解题思路】

4个木棒拼正三角形，就是2根一样长的和另外2根拼起来

木棒长度最大5000，可以用一个桶存下

枚举两根短的木棒的长度就可以了

【code】

#include <cstdio>
#include <iostream>
#include <algorithm>
using namespace std;
const int mod=1000000007;
int n,ans,Maxn=-1<<30;
int a[100005]; 
int cnt[5005];
inline int C1(int x){
    return x;
}
inline int C2(int x){
    return x*(x-1)/2;
}
int main(){
    scanf("%d",&n);
    for(register int i=1;i<=n;i++){
        scanf("%d",&a[i]);
        cnt[a[i]]++;
        if(a[i]>Maxn)Maxn=a[i];//桶计数
    }
    for(register int i=1;i<=Maxn;i++){
        if(cnt[i]>=2){
            int t=C2(cnt[i])%mod;
            for(register int j=1;j<=(i/2);j++){//枚举两根短木棒，剪枝
                if(j*2==i&&cnt[j]>=2)
                    ans=(ans+C2(cnt[j])*t)%mod; 
                else if(j*2!=i&&cnt[j]>=1&&cnt[i-j]>=1)
                    ans=(ans+C1(cnt[j])*t%mod*C1(cnt[i-j])%mod)%mod;//组合数计算
                }
        }
    }
    printf("%d\n",ans%mod);
    return 0;
}