2354. Number of Excellent Pairs

You are given a 0-indexed positive integer array nums and a positive integer k.

A pair of numbers (num1, num2) is called excellent if the following conditions are satisfied:

• Both the numbers num1 and num2 exist in the array nums.
• The sum of the number of set bits in num1 OR num2 and num1 AND num2 is greater than or equal to k, where OR is the bitwise OR operation and AND is the bitwise AND operation.

Return the number of distinct excellent pairs.

Two pairs (a, b) and (c, d) are considered distinct if either a != c or b != d. For example, (1, 2) and (2, 1) are distinct.

Note that a pair (num1, num2) such that num1 == num2 can also be excellent if you have at least one occurrence of num1 in the array.

• prefix sum solution

class Solution {
public:
    long long countExcellentPairs(vector<int>& A, int k) {
        unordered_map<int, int> pcnt;
        vector<int> psum(66);
        for(auto& a : A) pcnt[a] = __builtin_popcount(a);
        for(auto [_,v] : pcnt) psum[v] += 1;
        for(int i = 1; i < 66; i++) psum[i] += psum[i-1];

        long long res = 0;
        for(auto [me, cnt] : pcnt) {
            res += psum.back() - psum[max(k-cnt-1,0)];
        }

        return res;
    }
};

• fenwick tree solution

class Solution {
    int fenwick[1010] = {};
    void update(int n, int v) {
        while(n < 1010) {
            fenwick[n] += v;
            n += n & -n;
        }
    }
    int query(int n) {
        int res = 0;
        while(n) {
            res += fenwick[n];
            n -= n & -n;
        }
        return res;
    }
public:
    long long countExcellentPairs(vector<int>& A, int k) {
        unordered_map<int, int> pcnt;
        for(auto& a : A) {
            pcnt[a] = __builtin_popcount(a);
        }
        for(auto [_,v] : pcnt) {
            update(v, 1);
        }

        long long res = 0;
        for(auto [me, cnt] : pcnt) {
            res += query(1000) - query(max(k-cnt-1,0));
        }

        return res;
    }
};