LeetCode611-转化为桶+集合划分思想+生成函数

CSDN
考虑求出res[i]表示a[x]+a[y]=i,x!=y的下标对(x,y)个数。对res做前缀和，则得a[x]+a[y]<=i,x!=y的下标对(x,y)个数，记为ress。设b[x]表示a数组里等于x的元素个数（即开桶思想）。则ress[i]*b[i]表示最大值等于i且不能构成三角形的三元组个数，显然这是不合法的三元组集合的一个划分。则
$$ans=C_n^3-\sum _{1<=i<=m}ress[i]\ast b[i]$$
m表示a数组的最大值，n表示a数组非0元素个数。

求res是经典的生成函数问题。设b数组表示多项式B，答案与B^2有关。但这还不是答案，它包含了自己与自己组合的情况，并且每个二元组都算了2次。自己与自己组合，用B(x^2)表示，则所求多项式为：
$$\frac{B^2(x)-B(x^2)}{2}$$
用NTT求B^2，则复杂度为O(mlogm)。这里是暴力乘的。

class Solution {
#define rep(i,a,b) for(int i = (a);i <= (b);++i)
#define re_(i,a,b) for(int i = (a);i < (b);++i)
#define dwn(i,a,b) for(int i = (a);i >= (b);--i)
public:
    int triangleNumber(vector<int>& a) {
        int n = 0,m = 0;
        vector<int> b(1005);
        for(int x: a) if(x) n++,b[x]++,m = max(m,x);
        vector<int> res(2*m+5);
        rep(i,0,m) rep(j,0,i) res[i] += b[j]*b[i-j];
        rep(i,0,m) res[2*i] -= b[i];
        rep(i,0,m) res[i] /= 2;
        rep(i,1,m) res[i] += res[i-1];
        int ans = n*(n-1)*(n-2)/6;
        rep(i,0,m) ans -= b[i]*res[i];
        return ans;
    }
};