P3799 妖梦拼木棒

题目描述

有 n 根木棒，现在从中选 4 根，想要组成一个正三角形，问有几种选法？
答案对 (10^9+7) 取模。
因为固定选4根，所以，必定有两根是作为三角形边长，而另外两根的和为边长。
思路：

1. 统计一下每种长度的木棍个数
2. 枚举可能的三角形长度i
3. 枚举另外两个木棍中的一个的长度(j in [1, i - j])
4. if(j == i - j) res += ({2 choose cnt[j]} * {2 choose cnt[i]}) else res += ({1 choose cnt[j]} * {1 choose cnt[i - j]} * {2 choose cnt[i]})

这题有个大坑：就是在循环另外两个木棍其中一个的长度的时候必须写for(int j = 1; j <= i - j; j ++), 否则会出现同样的情况会计算两遍，结果就错了。

代码

#include<iostream>
#include<cstring>

using namespace std;

const int N = 100010, M = 5010, mod = 1e9 + 7;

#define LL long long

int n;
LL cnt[M];

LL calc(LL x){
    return x % mod * (x - 1) % mod / 2;
}

int main(){
    cin >> n;
    
    for(int i = 0; i < n; i ++){
        int x;
        cin >> x;
        
        cnt[x] ++;
    }
    
    LL res = 0;
    for(int i = 2; i <= 5000; i ++){
        if(cnt[i] < 2) continue;
        for(int j = 1; j <= i - j; j ++){ // 坑死我了
            LL t;
            if(j == i - j) t = calc(cnt[j]) % mod;
            else t = cnt[j] % mod * cnt[i - j] % mod;
            res = (res + calc(cnt[i]) % mod * t % mod) % mod;
        }
    }
    
    cout << res << endl;
}