1803

1803. Count Pairs With XOR in a Range (Hard)

Given a (0-indexed) integer array nums and two integers low and high, return the number of nice pairs.

A nice pair is a pair (i, j) where 0 <= i < j < nums.length and low <= (nums[i] XOR nums[j]) <= high.

 

Example 1:

Input: nums = [1,4,2,7], low = 2, high = 6
Output: 6
Explanation: All nice pairs (i, j) are as follows:
    - (0, 1): nums[0] XOR nums[1] = 5 
    - (0, 2): nums[0] XOR nums[2] = 3
    - (0, 3): nums[0] XOR nums[3] = 6
    - (1, 2): nums[1] XOR nums[2] = 6
    - (1, 3): nums[1] XOR nums[3] = 3
    - (2, 3): nums[2] XOR nums[3] = 5

Example 2:

Input: nums = [9,8,4,2,1], low = 5, high = 14
Output: 8
Explanation: All nice pairs (i, j) are as follows:
​​​​​    - (0, 2): nums[0] XOR nums[2] = 13
    - (0, 3): nums[0] XOR nums[3] = 11
    - (0, 4): nums[0] XOR nums[4] = 8
    - (1, 2): nums[1] XOR nums[2] = 12
    - (1, 3): nums[1] XOR nums[3] = 10
    - (1, 4): nums[1] XOR nums[4] = 9
    - (2, 3): nums[2] XOR nums[3] = 6
    - (2, 4): nums[2] XOR nums[4] = 5

 

Constraints:

• 1 <= nums.length <= 2 * 104
• 1 <= nums[i] <= 2 * 104
• 1 <= low <= high <= 2 * 104

Companies: Google, Vimeo

Related Topics:
Array, Bit Manipulation, Trie

Hints:

• Let's note that we can count all pairs with XOR ≤ K, so the answer would be to subtract the number of pairs withs XOR < low from the number of pairs with XOR ≤ high.
• For each value, find out the number of values when you XOR it with the result is ≤ K using a trie.

Solution 1. Trie

Complexity Analysis

Adding a number into trie takes O(16) = O(1) time.

The count will at most visit a perfect binary-tree like trie whose height is 16, so there are at most O(2^16-1) = O(1) nodes. Hence the space complexity is O(1).

Since we repeat the above operations N times, the overall time complexity is O(N).

// OJ: https://leetcode.com/problems/count-pairs-with-xor-in-a-range/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(1)
struct TrieNode {
    TrieNode *next[2] = {};
    int cnt = 0;
};
const int BIT = 15;
class Solution {
    void addNode(TrieNode *node, int n) {
        for (int i = BIT; i >= 0; --i) {
            int b = n >> i & 1;
            if (!node->next[b]) node->next[b] = new TrieNode();
            node = node->next[b];
            node->cnt++;
        }
    }
    int count(TrieNode *node, int n, int low, int high, int i = BIT, int rangeMin = 0, int rangeMax = (1 << (BIT + 1)) - 1) {
        if (rangeMin >= low && rangeMax <= high) return node->cnt;
        if (rangeMax < low || rangeMin > high) return 0;
        int ans = 0, b = n >> i & 1, r = 1 - b, mask = 1 << i;
        if (node->next[b]) ans += count(node->next[b], n, low, high, i - 1, rangeMin & ~mask, rangeMax & ~mask);
        if (node->next[r]) ans += count(node->next[r], n, low, high, i - 1, rangeMin | mask, rangeMax | mask);
        return ans;
    }
public:
    int countPairs(vector<int>& A, int low, int high) {
        TrieNode root;
        int ans = 0;
        for (int n : A) {
            ans += count(&root, n, low, high);
            addNode(&root, n);
        }
        return ans;
    }
};

Solution 2. Trie

// OJ: https://leetcode.com/problems/count-pairs-with-xor-in-a-range/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(1)
struct TrieNode {
    TrieNode *next[2] = {};
    int cnt = 0;
};
const int BIT = 15;
class Solution {
    void addNode(TrieNode *node, int n) {
        for (int i = BIT; i >= 0; --i) {
            int b = n >> i & 1;
            if (!node->next[b]) node->next[b] = new TrieNode();
            node = node->next[b];
            node->cnt++;
        }
    }
    int countLessThan(TrieNode *node, int n, int limit) {
        int ans = 0;
        for (int i = BIT; i >= 0 && node; --i) {
            int b = n >> i & 1, r = 1 - b, lb = limit >> i & 1;
            if (lb == 1) {
                if (node->next[b]) ans += node->next[b]->cnt;
                node = node->next[r];
            } else {
                node = node->next[b];
            }
        }
        return ans;
    }
    int count(TrieNode *node, int n, int low, int high) {
        return countLessThan(node, n, high + 1) - countLessThan(node, n, low);
    }
public:
    int countPairs(vector<int>& A, int low, int high) {
        TrieNode root;
        int ans = 0;
        for (int n : A) {
            ans += count(&root, n, low, high);
            addNode(&root, n);
        }
        return ans;
    }
};