三个数组中的三数之和求索引组合总数

题目

给定三个数组A，B，C，数组无序，且数组中元素都是正整数。求 A[i] + B[j] + C[k] = w的(i , j, k )的组合总数。要求时间O(N), 空间O(1)。

思路

对于每个数组建立一个桶（哈希map），每个数组中大于w的数直接忽略，其它数放入桶中，map的第一个元素是数值，第二个元素是频次，即该数值出现的次数。因此三个桶的大小都不大于64，所以占用空间与N无关，空间复杂度O(1)，接下来直接遍历统计有多少种相加等于w的组合，方式是让每个数对应的频次相乘即可。由于桶的大小是固定的，不大于64，与N无关，因此遍历时候的时间复杂度是O（1），而建桶的时间复杂度是O(N)，因此总的时间复杂度是O（N）。

代码

#include <bits/stdc++.h>
using namespace std;

int total(vector<int> &A, vector<int> &B, vector<int> &C, int k){
    int len_A = A.size();
    int len_B = B.size();
    int len_C = C.size();

    unordered_map<int, int> bucket_A;
    unordered_map<int, int> bucket_B;
    unordered_map<int, int> bucket_C;

    for(auto &a : A){
        if(a < 64) bucket_A[a]++;
    }
    for(auto &b : B){
        if(b < 64) bucket_B[b]++;
    }
    for(auto &c : C){
        if(c < 64) bucket_C[c]++;
    }

    int res = 0;
    for(auto &a : bucket_A){
        for(auto &b : bucket_B){
            if(a.first + b.first < 64){
                auto c = bucket_C.find(k - a.first - b.first);
                if(c != bucket_C.end()){
                    res += a.second * b.second * c->second;
                }
            }
        }
    }

    return res;
}

int main(){
    vector<int> A{1, 2, 3, 4, 5, 65, 1};
    vector<int> B{1, 2, 3, 4, 5};
    vector<int> C{1, 2, 3, 4, 5};
    
    int res = total(A, B, C, 15);
    cout<<res<<endl;
    return 0;

}