c_cpp 给定一个整数数组，找到三元组，使得a ^ 2 + b ^ 2 = c ^ 2。

Posted 2021-05-14

vector<vector<int>> find_triplets(const vector<int>& A) {
    vector<vector<int>> res;
    int N = A.size();
    if(N < 3) return res;
    
    vector<long long> B;
    for(int i : A) B.push_back(i*i); // gist! push the square values
    
    // sort them in descending order not ascending order! 
    sort(B.begin(), B.end(), greater<int>());
    
    for(int i=0; i<N-2; ++i) {
        int j = i+1, k = N-1;       // i+1 not i
        while(j < k) {              // < not <=
            if(B[i] == B[j] + B[k]) {
                res.push_back({A[i], A[j], A[k]});
                ++j;                // forgot!!
                --k;
            }
            else if(B[i] < B[j] + B[k]) j++;
            else k--;
        }
    }
    return res;
}