P3799 妖梦拼木棒

题目链接：

欲由4根木棒组成一个正三角形，则必有2根长度相等，且另外2根长度之和，等于前2根相等的木棒的长度。

由于各木棍的长度 $a_i⩽5000$，时间复杂度 $O(n^2)$ 可过。考虑直接用两层循环，暴力枚举上述两种木棒的长度，计算方案数并累加。

外层循环为从长度为 $i$ 的木棍中取出 $2$ 根，方案数为 $\mathrm{C}(\mathrm{c}\mathrm{n}\mathrm{t}[\mathrm{i}],2)$。

内层循环为从剩余的木棍中取出两根长度之和为 $i$ 的木棍。设其中一根长度为 $j$， 则另外一根长度为 $i-j$。

①若 $j=i-j$，则方案数为从 $\mathrm{c}\mathrm{n}\mathrm{t}[\mathrm{j}]$ 中取出2个数的组合，即$\mathrm{C}(\mathrm{c}\mathrm{n}\mathrm{t}[\mathrm{j}],2)$

②若 $j\ne i-j$，需从长度为 $j$ 和 $i-j$ 的木棍中各取出一个，由乘法原理知方案数为 $\mathrm{C}(\mathrm{c}\mathrm{n}\mathrm{t}[\mathrm{j}],2)\cdot \mathrm{C}(\mathrm{c}\mathrm{n}\mathrm{t}[\mathrm{i}-\mathrm{j}],2)$

一定注意随时取模。

#include <cstdio>
#include <algorithm>
 
using LL = long long;
 
const int N = 1e5 + 5, mod = 1e9 + 7;
 
int n, a[N], cnt[N];//cnt[i]表示长度为i的木棍个数
 
int main()
{
	scanf("%d", &n);
	int min = 1e9, max = -1e9;
	for (int i = 0; i < n; i++) {
		scanf("%d", &a[i]);
		cnt[a[i]]++;
		min = std::min(min, a[i]);
		max = std::max(max, a[i]);
	}
	
	LL ans = 0;
	for (int i = min; i <= max; i++) {
		if (cnt[i] >= 2) {//为保证可构成三角形，此长度的木棒数量 ≥2时才可进入内层循环。
		    LL times = cnt[i] * (cnt[i] - 1) / 2 % mod;
			for (int j = min; j <= i / 2; j++) {//为避免重复计算，规定j <= i - j
				if (j != i - j && cnt[j] >= 1 && cnt[i - j] >= 1) ans += times * cnt[j] * cnt[i - j] % mod;
				else if (j == i - j && cnt[j] >= 2) ans += times * (cnt[j] * (cnt[j] - 1) / 2) % mod;
			}
		}
	}
	printf("%lld", ans%mod);
	return 0;
}