洛谷P3799 妖梦拼木棒

妖梦拼木棒

题目背景

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

题目描述

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

答案对 $10^9+7$ 取模。

输入格式

第一行一个整数 $n$。

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

输出格式

一行一个整数代表答案。

样例 #1

样例输入 #1

4 
1
1
2
2

样例输出 #1

1

提示

数据规模与约定

• 对于 $30\%$ 的数据，保证 $n\le 5\times 10^3$。
• 对于 $100\%$ 的数据，保证 $1\le n\le 10^5$，$1\le a_i\le 5\times 10^3$。

题解

四根木棍, 拼成一个正三角形. 显而易见的是: 一定有两根是一样长的, 另外两根加起来也与其余的一样长.
那么我们就可以用两层循环来做这道题
外层循环用来找两根一样长的,内层循环用来找那两根短一点的.

在代码阶段我们需要考虑题目中数据要怎么去存储
对于数据

1
1
2
2

保存为data=[1,1,2,2]显然不是一种好方法. 因为你在后续从中找数据的时候就需要遍历里面所有元素
所以采用data[1]=2, data[2]=2数组中的值就是该长度木根的个数
这样以后, 在我们找那两根短一点的木根的时候优势就显现出来了.
对于边长为a的正三角形, 那两根短一些的木棍, 当选定了一根的长度为m, 另一根的长度就是a-m, 这时候只需要到数组中去找对应的值是否不为0就可以了

#include <stdio.h>
#define ll long long
const int kMaxn = 1e6 + 10;
const ll kMod = 1e9 + 7;

ll C(ll x, ll k) {
    //求得从n个数中取出k个数的组合
    //此处k=1 / 2，用了特判写法。
    //k = 1 时，C(x, 1) = x;
    //k = 2 时，C(x, 2) = x * (x - 1) / 2;
    return (k == 1ll ? x : x * (x - 1ll) / 2ll) % kMod;
}

int main() {
    ll n, ans, maxa;
    ll a[kMaxn], num[kMaxn];
    scanf("%lld\n", &n);

    maxa = 0;
    for (int i = 1; i <= n; ++i) { //读入,并放入计数数组中 
        scanf("%lld", &a[i]);
        maxa = (a[i] > maxa) ? a[i] : maxa;
        num[a[i]]++;
    }

    ans = 0;
    for (int i = 2; i <= maxa; ++i) { //枚举两根相等的边
        if (num[i] >= 2ll) {
            ll times = C(num[i], 2ll) % kMod; //求出组合数 
            for (int j = 1; j <= i / 2; ++j) { //枚举被合成的边(到i / 2即可)
                if (j != i - j && num[j] >= 1 && num[i - j] >= 1) //用来合成的木棒长度不等 
                    ans += times * C(num[j], 1) * C(num[i - j], 1) % kMod;
                if (j == i - j && num[j] >= 2) //用来合成的木棒长度相等 
                    ans += times * C(num[j], 2) % kMod;
                ans %= kMod;
            }
        }
    }

    printf("%lld", ans);
    return 0;
}