洛谷题单指南-暴力枚举-P3799 妖梦拼木棒

原题链接：https://www.luogu.com.cn/problem/P3799

题意解读：要选四根木棒拼成等边三角形，必然有两根长度相等，其余两根长度之和等于前两根

解题思路：

木棒总数最大100000，每根最长5000，因此通过枚举其中两根木棒的长度，计算出另外两根的长度，通过各个长度的木棒数进行选择。

设数组h[n]保存每种长度木棒的数量, 设四根木棍长度为a < b < c = d

枚举a，b的长度, a从1~5000，b从a~5000，且需要符合a + b <= 5000

如果a == b，则在长度为a的木棒中选2个，即C(h[a], 2)，再在长度为a + b的木棒中选2个，即C(h[a+b], 2)，

根据乘法原理，方案数为C(h[a], 2) * C(h[a+b], 2)

如果a != b，则在长度为a的木棒中选1个，即C(h[a], 1)，在长度为b的木棒中选1个，即C(h[b], 1)，再在长度为a + b的木棒中选2个，即C(h[a+b], 2)，

根据乘法原理，方案数为C(h[a], 1) * C(h[b], 1) * C(h[a+b], 2)

注意每一次计算中需要对1e9 + 7取模。

100分代码：

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

const int N = 5005;
const int MOD = 1e9 + 7;

long long h[N];
long long ans;
int n;

long long C(long long a, int b)
{
    if(b == 1) return a;   // 计算C(a, 1)
    else return a * (a - 1) / 2 % MOD; // 计算C(a, 2)
}

int main()
{
    cin >> n;
    int x;
    for(int i = 1; i <= n; i++)
    {
        cin >> x;
        h[x]++;
    }

    for(int i = 1; i <= 5000; i++) //枚举第一根木棒，设定第一根最短
    {
        for(int j = i; i + j <= 5000; j++) //枚举第二根，第二根大于等于第一根，且两者之后不超过最大长度
        {
            if(h[i + j] < 2) continue;

            if(i != j && h[i] && h[j]) //选一个i，选一个j，选两个i+j
            {
                ans += C(h[i], 1) * C(h[j], 1) * C(h[i + j], 2) % MOD;
            }

            if(i == j && h[i] >= 2) //选两个i，选两个i+j
            {
                ans += C(h[i], 2) * C(h[i + j], 2) % MOD;
            }
        }
    }

    cout << ans % MOD;

    return 0;
}