P3799 妖梦拼木棒

题目传送门

#include <bits/stdc++.h>

using namespace std;
typedef long long LL;
const LL MOD = 1e9 + 7;//要模的质数

const int N = 1e5 + 10;
int a[N]; //木棍长度数组,木棍数量最大值是10^5,所以N需要开到1e5+10
int b[N]; //每个长度的木棍出现的次数

int n;      //n根木棍
int MAX;    //木棍的最大长度
LL res;     //结果,为防止在运算过程中暴INT，声明成LL


//组合数迭代模板
int c[N][4]; //为什么要声明N*5维的数组呢？第一维是说，最多可能需要从N个元素中取，第二维是说，最多取4个.为啥是4个呢？因为题目说取四根最多了
void init() {
    for (int i = 0; i < N; i++)
        for (int j = 0; j <= min(i, 4); j++)
            if (!j)c[i][j] = 1;
            else c[i][j] = (c[i - 1][j] + c[i - 1][j - 1]) % MOD;
}

int main() {
    //生成数据范围内的所有组合数
    init();

    //n根木棒
    cin >> n;
    //输入数据
    for (int i = 1; i <= n; ++i) {
        cin >> a[i];
        //1、记录输入的最大值
        MAX = max(a[i], MAX);
        //2、使用桶记录每个长度出现的次数
        b[a[i]]++;
    }

    //遍历每一个可能存在的长度
    for (int i = 2; i <= MAX; i++) {
        //只有出现两次及以上的边才有机会成为长边
        if (b[i] >= 2) {
            //在b[i]中取2个有多少种方法，记得每步取模
            LL cnt = c[b[i]][2];

            //要从剩余的木棒中取出2根长度之和为i的木棒。令其中一根长度为j，则另一根长度为i-j。
            for (int j = 1; j <= i / 2; j++) {
                //如果两根一样长，那么就需要这个小的最小有2根
                //另一根就是i-j的长度
                if (i == 2 * j && b[j] >= 2)
                    res += cnt * c[b[j]][2] % MOD;

                //两种情况之间是加法原理
                //用来合成的木棒长度不等
                if (i != 2 * j && b[j] >= 1 && b[i - j] >= 1)
                    res += cnt * c[b[j]][1] * c[b[i - j]][1] % MOD;

                //一定注意随时取模
                res %= MOD;
            }
        }
    }
    printf("%lld", res);
    return 0;
}