题解：洛谷 P3799 小 Y 拼木棒

【题目来源】

洛谷：P3799 小 Y 拼木棒 - 洛谷 (luogu.com.cn)

【题目描述】

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

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

答案对 \(10^9+7\) 取模。

【输入】

第一行一个整数 \(n\)。

第二行往下 \(n\) 行，每行 \(1\) 个整数，第 \(i\) 个整数 \(a_i\) 代表第 \(i\) 根木棒的长度。

【输出】

一行一个整数代表答案。

【输入样例】

4 
1
1
2
2

【输出样例】

1

【解题思路】

【算法标签】

《洛谷 P3799 妖梦拼木棒》 #枚举# #组合数学# #洛谷原创#

【代码详解】

#include <bits/stdc++.h>
using namespace std;

int main()
{
    // counts数组用于统计每个数字出现的次数，初始化为0
    long long counts[5005] = {0}, n, tmp, ans = 0;
    cin >> n; // 输入数字的总个数n

    // 读取n个数字，并统计每个数字出现的次数
    for (int i = 1; i <= n; i++) 
	{
        cin >> tmp; // 输入当前数字
        counts[tmp]++; // 对应数字的计数加1
    }

    // 遍历所有可能的i和j，计算满足条件的组合数
    for (int i = 1; i <= 2500; i++) 
	{
        for (int j = i; j <= 5000 - i; j++) 
		{
            // k3表示数字i+j的出现次数
            int k3 = counts[i + j];
            // k1表示数字i的出现次数
            int k1 = counts[i];
            // k2表示数字j的出现次数
            int k2 = counts[j];

            // 如果i+j至少出现2次，才可能有有效组合
            if (k3 >= 2) 
			{
                // 情况1：i == j，即两个相同的数相加
                if (i == j && k1 >= 2) 
				{
                    // 组合数计算：C(k1,2) * C(k3,2)
                    tmp = (k1 * (k1 - 1) / 2) * (k3 * (k3 - 1) / 2) % 1000000007;
                    ans += tmp;
                }
                // 情况2：i != j，即两个不同的数相加
                if (i != j && k1 >= 1 && k2 >= 1) 
				{
                    // 组合数计算：k1 * k2 * C(k3,2)
                    tmp = k1 * k2 * (k3 * (k3 - 1) / 2) % 1000000007;
                    ans += tmp;
                }
            }
        }
    }

    // 对结果取模，防止溢出
    ans %= 1000000007;
    cout << ans; // 输出最终结果
    return 0;
}

【运行结果】

4 
1
1
2
2
1