洛谷 P3799 妖梦拼木棒

题目
这道题主要考组合数学，一开始我没注意数据范围，傻乎乎地模拟，结果复杂度爆炸10个TLE。这道题是算出来的。
思路是用一个栈把相同长度个数超过2的木棍用栈保存下来，同时用一个数组存下所有长度木棍的个数。计数时从栈里拿出一个个数超过2的木棍A，通过组合公式算出从中选出2个的选法，再从剩下的木棍里找出两个木棍B，C，使A = B + C；四根木棍分别是A A B C；

#include <iostream>
#include <cmath>
#include <iomanip>
#include <vector>
#include <stack>
#include <cstring>
#include <algorithm>
using namespace std;

typedef long long ll;
const int Mod = 1e9 + 7;
stack<int> stck;
int stick[5001];//存放该长度的个数
ll n, tot;

int C(int N, int M)//组合公式，题目里只有两种情况，选一个或者选两个
{
    if (M == 1) {// 选一个
         return stick[N];
    }
    else {  //选两个
        return stick[N] * (stick[N] - 1) / 2;
    }
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);

    int n;
    cin >> n;
    int tmp;
    for (int i = 0; i < n; i++) {
        cin >> tmp;
        stick[tmp]++;
        if (stick[tmp] == 2) {
            stck.push(tmp);//入栈
        }
    }
    while (stck.size()) {
        int tmp = stck.top();
        stck.pop();
        int A = C(tmp, 2);//2根长木棍的选法
        int l = tmp / 2, r = tmp / 2; //从中间往两边找
        if (tmp % 2) {//如果长木棍的长度是奇数右边木棍的长度+1
            r++;	  //例如tmp = 5, l = 2, r = 3;
        }
        int cnt = 0;
        for (int i = l, j = r; i > 0 && j < tmp; i--, j++) {//两根短木棍，较短的不能小于等于0
            if (i + j == tmp) {								//不能大于等于长木棍		
                if (i == j) {//两根木棍同
                    cnt += C(i, 2);
                }
                else {
                    cnt += C(i, 1) * C(j, 1);
                }
            }
        }
        tot += A * cnt;//记录总选法
    }
    cout << tot % Mod << endl;
    return 0;
}