P3799 妖梦拼木棒【暴力枚举、组合数】

题目背景

上道题中，妖梦斩了一地的木棒，现在她想要将木棒拼起来。

题目描述

有 nn根木棒，现在从中选 4 根，想要组成一个正三角形，问有几种选法？

答案对 10^9+7 取模。

输入格式

第一行一个整数 nn。

第二行 nn 个整数，第 i 个整数 ai​ 代表第 i 根木棒的长度。

输出格式

一行一个整数代表答案。

输入输出样例

输入

4 
1 1 2 2

输出

1

说明/提示

数据规模与约定

• 对于 30% 的数据，保证 n≤5×10^3。
• 对于 100% 的数据，保证 1≤n≤10^5，0≤ai​≤5×10^3。

解题思路

首先，我们分析题意，可知要选择四根木棍组成正三角形，可知一定有两根木根作为完整的一条边且长度相等。另外，有两根木棍通过合并组成另一条边。这里我们假设三角形的边长为i，那么有两条长度为i的边，那么另两条较短的木棍长度为j和i-j。

我们先用一个num[]数组存储所有的木棍的数量，同时利用max和min记录所有木棍中最长和最短的木棍。

然后通过i从min+1到max遍历所有木棍（不用从min开始，因为j要更小些），然后j从min到i/2遍历。其中注意，如果num[i]<2那么就无需再进行内层的j的遍历了，因为i的数量不够。如果num[i]>=2，那么我们进行j的遍历，其中有两种情况：

1. 当j!=i-j，ans=ans+C(num[i]，2)*C(num[j]，1)*C(num[i-j]，1)%mod;
2. 当j==i-j且num[j]<=2时，ans=ans+C(num[i]，2)*C(num[j]，2)%mod;

其中C(n，m)表示组合数。

代码 

#include<cstdio>
#include<algorithm>
using namespace std;
#define ll long long
#define MAXN 100010
const int  mod = 1e9 + 7;

int n;
int num[MAXN] = { 0 };

int main()
{
	ll ans = 0;
	scanf("%d", &n);
	int max = 0;
	int min = MAXN;
	for (int i = 0; i < n; i++)
	{
		int a;
		scanf("%d", &a);
		num[a]++;
		if (a > max)
			max = a;
		if (min > a)
			min = a;
	}
	for (int i = min+1; i <= max; i++)    //i得比j大
	{
		if (num[i] < 2)
			continue;
		for (int j = min; j <= i / 2; j++)    //j得比i小
		{
			if (num[j] && num[i - j])  //如果长度为j和长度为i-j的木棒存在
			{
				if (j != i - j)
				{
					ans=ans+((num[i]*(num[i]-1)>>1)*num[j]*num[i-j])%mod;
				}
				else if (num[j]>=2)
				{
					ans = ans + ((num[i] * (num[i] - 1) >> 1)*(num[j] * (num[j] - 1) >> 1) % mod) % mod;
				}
			}
		}
	}
	printf("%lld\n", ans%mod);
	return 0;
}