LeetCode 923. 3Sum With Multiplicity

Quora的OA Forming Team，其实就是Three Sum的变种，但是麻烦不少。

方法一：

老老实实Two Pointer，固定第一个（重复不跳），第二个第三个比较复杂。如果a[low]==a[high]，有C_n^2种可能，如果不相等，那就要向右数和a[low]相等的元素个数*向左数和a[high]相等的个数。

class Solution {
public:
    int threeSumMulti(vector<int>& A, int target) {
        int MOD=int(1e9+7);
        sort(A.begin(),A.end());
        int count=0;
        for (int i=0;i<A.size();++i){
            int low=i+1, high=A.size()-1;
            while (low<high){
                int sum=A[i]+A[low]+A[high];
                if (sum<target) ++low;
                else if (sum>target) --high;
                else{
                    if (A[low]==A[high]){
                        int tmp=high-low+1;
                        count = (count+tmp*(tmp-1)/2)%MOD;
                        //cout << i << ' ' << low << ' ' << high << ' ' << tmp*(tmp-1)/2 << endl;
                        break;
                    }else{
                        //cout << i << ' ' <<  low << ' ' << high << ' ';
                        ++low; --high;
                        int num1=1, num2=1;
                        while (low<=high && A[low]==A[low-1]){++num1; ++low;}
                        while (low<=high && A[high]==A[high+1]){++num2; --high;}
                        count = (count+num1*num2)%MOD;
                        //cout << num1*num2 << endl;
                    }
                }
            }
        }
        return count;
    }
};

 

方法二：

先计算每个元素的个数，然后按照不能重复的3Sum来做就行了。同样需要注意各种情况，但是思路比上一种方法更清晰一点。

class Solution {
public:
    int threeSumMulti(vector<int>& a, int target) {
        sort(a.begin(),a.end());
        long res=0;
        
        unordered_map<int,long> count; // use long for convenience
        for (auto x:a) ++count[x];
        
        for (int i=0;i<a.size();++i){
            if (i>0 && a[i]==a[i-1]) continue;
            int low=i+1, high=a.size()-1;
            while (low<high){
                int sum=a[i]+a[low]+a[high];
                if (sum<target) ++low;
                else if (sum>target) --high;
                else{
                    if (a[i]!=a[low] && a[low]!=a[high]){ // all different
                        res += count[a[i]]*count[a[low]]*count[a[high]];
                    }else if (a[i]==a[low] && a[low]!=a[high]){
                        res += count[a[i]]*(count[a[i]]-1)/2 * count[a[high]];
                    }else if (a[i]!=a[low] && a[low]==a[high]){
                        res += count[a[i]] * count[a[low]]*(count[a[low]]-1)/2;
                    }else{ // all same  C_n^3
                        res += count[a[i]]*(count[a[i]]-1)*(count[a[i]]-2)/6;
                    }
                    
                    ++low; --high;
                    while (low<high && a[low]==a[low-1]) ++low;
                    while (low<high && a[high]==a[high+1]) --high;
                }
            }
        }
        return res % int(1e9+7);
    }
};