洛谷P3799 妖夢拼木棒

題目背景

上道題中，妖夢斬了一地的木棒，現在她想要將木棒拼起來。

題目描述

有n根木棒，現在從中選4根，想要組成一個正三角形，問有幾種選法？

輸入輸出格式

輸入格式：

第一行一個整數n

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

輸出格式：

一行一個整數，對1e9+7取模

輸入輸出樣例

4 1 1 2 2

1

說明

對於30%的數據 N<=5000

對於100%的數據 N<=100000

by-szc

果然數學還是自己的短板啊（其實自己什麼都很垃圾）

這道題應該是比較裸的組合數問題

枚舉短的邊，查詢長的邊

組合數預先處理

// luogu-judger-enable-o2
#include<iostream>
#include<cstdio>
#define LL long long 
using namespace std;
const LL MAXN=2*1e6+10;
const LL mod=1e9+7;
LL N;
LL a[MAXN],h[MAXN];
LL C1(LL a){return a;}
LL C2(LL a){return ( a*(a-1) )/2; }
int main()
{
    #ifdef WIN32
    freopen("a.in","r",stdin);
    #else 
    #endif
    LL ans=0;
    scanf("%lld",&N);
    for(LL i=1;i<=N;i++) scanf("%lld",&a[i]),h[ a[i] ] ++;
    for(LL i=1;i<=5000;i++)
        for(LL j=i;j<=5000&&i+j<=5000;j++)
        {
            if(i==j) 
            {
                if(h[i]>=2&&h[i*2]>=2) ans+=C2(h[i])*C2(h[i*2])%mod; 
            }
            else     
            {
                 if(h[i]>=1&&h[j]>=1&&h[i+j]>=2) 
                     ans+=C1(h[i])*C1(h[j])*C2(h[i+j])%mod;
            } 
        }
            
    printf("%lld",ans%mod);
    return 0;
}