力扣OJ 454. 四数相加 II

给定四个包含整数的数组列表 A , B , C , D ,计算有多少个元组 (i, j, k, l) ，使得 A[i] + B[j] + C[k] + D[l] = 0。

为了使问题简单化，所有的 A, B, C, D 具有相同的长度 N，且 0 ≤ N ≤ 500 。所有整数的范围在 -228 到 228 - 1 之间，最终结果不会超过 231 - 1 。

例如:

输入:
A = [ 1, 2]
B = [-2,-1]
C = [-1, 2]
D = [ 0, 2]

输出:
2

解释:
两个元组如下:
1. (0, 0, 0, 1) -> A[0] + B[0] + C[0] + D[1] = 1 + (-2) + (-1) + 2 = 0
2. (1, 1, 0, 0) -> A[1] + B[1] + C[0] + D[0] = 2 + (-1) + (-1) + 0 = 0

//给vector拓展，加上id并排序
template<typename T>
bool cmp(T x,T y)
{
    return x<y;
}
template<typename T>
vector<pair<T,int>> sortWithId(vector<T>v)
{
    vector<pair<T,int>>ans;
    ans.resize(v.size());
    for(int i=0;i<v.size();i++)ans[i].first=v[i],ans[i].second=i;
    sort(ans.begin(),ans.end(),[](pair<T,int>p1,pair<T,int>p2){return cmp(p1.first,p2.first);});
    return ans;      
}
//2个vector中寻找和为s的对，返回结果的每一行都是[id1,id2]
template<typename T>
vector<vector<int>> findSum(vector<T>v1,vector<T>v2,T s)
{
    vector<vector<int>>ans;
    int m=min((long long)(v1.size()*v2.size()),(long long)12345678);
    ans.reserve(m);
    vector<int>tmp(2);
    vector<pair<T,int>>v3=sortWithId(v2);
    sort(v2.begin(), v2.end(),cmp<T>);
    for(int i=0;i<v1.size();i++)
    {
        auto it1=lower_bound(v2.begin(), v2.end(), s-v1[i]);
        auto it2=upper_bound(v2.begin(), v2.end(), s-v1[i]);
        tmp[0]=i;
        for(auto j=it1;j<it2;j++)
        {
            tmp[1]=v3[j-v2.begin()].second;
            ans.push_back(tmp);
        }
    }
    return ans;
}
//枚举2个vector的两数之和
template<typename T>
vector<T> everySum(vector<T>&v1,vector<T>&v2)
{
    vector<T>ans;
    ans.resize(v1.size()*v2.size());
    int k=0;
    for(int i=0;i<v1.size();i++)for(int j=0;j<v2.size();j++)ans[k++]=v1[i]+v2[j];
    return ans;
}
//判断有序数组中有多少个不同的数
template<typename T>
int getDifNum(vector<T>v)
{
    int ans=1;
    for(int i=1;i<v.size();i++)ans+=(v[i]!=v[i-1]);
    return ans;
}
//收缩计数,把[1 1 4 4 4]变成[(1,2)(4,3)]
template<typename T>
vector<pair<T,int>> fshr(vector<T>v)
{
    vector<pair<T,int>>ans;
    if(v.size()==0)return ans;
    sort(v.begin(),v.end());
    ans.resize(getDifNum(v));
    ans[0]=make_pair(v[0],1);
    for(int i=1,j=0;i<v.size();i++)
    {
        if(v[i]==v[i-1])ans[j].second++;
        else ans[++j]=make_pair(v[i],1);
    }
    return ans;
}
//提取pair数组的first
template<typename T1,typename T2>
vector<T1> fdraw(vector<pair<T1,T2>>v)
{
    vector<T1>ans(v.size());
    for(int i=0;i<v.size();i++)ans[i]=v[i].first;
    return ans;
}
//提取pair数组的second
template<typename T1,typename T2>
vector<T2> fdraw2(vector<pair<T1,T2>>v)
{
    vector<T2>ans(v.size());
    for(int i=0;i<v.size();i++)ans[i]=v[i].second;
    return ans;
}
class Solution {
public:
    int fourSumCount(vector<int>& A, vector<int>& B, vector<int>& C, vector<int>& D) {
        vector<int> v1 = everySum(A,B),v2 = everySum(C,D);
        vector<pair<int,int>>v3=fshr(v1),v4=fshr(v2);
        vector<int> v5=fdraw(v3),v6=fdraw(v4);
        vector<int> v7=fdraw2(v3),v8=fdraw2(v4);
        vector<vector<int>> ans = findSum(v5,v6,0);
        int res=0;
        for(int i=0;i<ans.size();i++)res+=v7[ans[i][0]]*v8[ans[i][1]];
        return res;
    }
};