P3799 妖梦拼木棒

题目背景

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

题目描述

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

答案对 10^9+7 取模。

输入格式

第一行一个整数 n。

第二行 n 个整数，第 i 个整数 a_i 代表第 i 根木棒的长度。

输出格式

一行一个整数代表答案。

思路

选择4根木棒组成一个正三角形，说明有两个长度相等的木棒构成三角形的两条边，另外两根木棒共同组成三角形的一个边，其长度和前面的木棒的长度相等。

代码

#include <iostream>
#include <algorithm>
#define ll long long
using namespace std;

const int maxn = 5005;
const ll mod = 1e9 + 7;

int num[maxn];  //num[i]表示长度为i的木棒的个数
int maxLen; //最长的木棒的长度
ll ans; //保存答案

//组合数公式，s为总数，取k个数（在这里只是1或者2）
ll C(ll s, ll k)
{
	return k == 1LL ? s : (s - 1) * s / 2;
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);

	int n;  //木棒个数
	cin >> n;

	int len;
	for (int i = 1; i <= n; i++)
	{
		cin >> len;
		num[len]++;
		maxLen = max(maxLen, len);
	}

	//枚举作为两条边的木棒的长度，由于第三条边由两根木棒组成，而木棒长度都是整数，所以从2开始枚举
	for (int i = 2; i <= maxLen; i++)
	{
		if (num[i] >= 2)
		{
			ll temp = C(num[i], 2);
			for (int j = 1; j <= i / 2; j++) //j为组成第三条边的较短木棒的长度
			{
				if (j != i - j && num[j] >= 1 && num[i - j] >= 1)  //如果两根木棒长度不相等
				{
					ans += temp * C(num[j], 1) * C(num[i - j], 1) % mod;
				}
				else if (j == i - j && num[j] >= 2)  //如果两根木棒的长度相等
				{
					ans += temp * C(num[j], 2) % mod;
				}
				ans %= mod;
			}

		}
	}
	cout << ans;
}