1803. Count Pairs With XOR in a Range

HardBit ManipulationTrieArray
Leetcode Link

Problem Description

In this problem, we're given an integer array nums and two integers low and high. Our task is to find the number of 'nice pairs' in the array. A 'nice pair' is defined as a pair of indices (i, j) such that 0 <= i < j < nums.length and low <= (nums[i] XOR nums[j]) <= high. The XOR here is the bitwise exclusive OR operation, which compares the bits of two numbers and returns 1 for each position where the corresponding bits of the two numbers are different, and returns 0 where they are the same.

Intuition

This problem is a good candidate for a Trie (prefix tree), especially when dealing with bits and XOR operations. A Trie allows us to efficiently store and retrieve binary representations of numbers. The main intuition is that, for each number x in nums, we can insert its binary representation into the Trie and simultaneously query the Trie to count how many numbers previously inserted into the Trie would form a 'nice pair' with x.

We need to calculate two quantities for each number x in nums:

  1. The count of numbers in Trie so far that form a 'nice pair' with x when XORed, such that the result is less than or equal to high.
  2. The count of numbers in Trie so far that form a 'nice pair' with x when XORed, such that the result is less than low.

The difference between these two counts gives us the number of 'nice pairs' for that particular value of x. By subtracting the count for low from the count for high + 1, we exclude pairs where the XOR is too small.

The Trie structure is a special tree where each node represents a bit position (from the most significant bit to the least), and each path from the root to a node represents the prefix of the binary representation of the inserted numbers. Every node stores the count of elements that share the same bit prefix up to that node. To ensure that we are considering all 16-bits of the integers, we iterate from the 15th to the 0th bit (since array elements are within the range of 0 to 10^4, and 10^4 in binary takes up to 14 bits, we use 16 for a safe bound).

The search method on the Trie tries to maximize the XOR result while keeping it under the given limit (high or low). At every bit, we decide whether to continue with the same bit or switch to the opposite bit to maximize the XOR based on the limit's current bit. If we can afford to switch the bit (based on the limit), we add the count of nodes under the current bit to the overall count and then move to the opposite bit for higher XOR value. Otherwise, we just follow the current bit path. This search method allows us to efficiently count all the suitable pairs for each entered number.

Adding all the differences together for each number in the array gives us the total count of 'nice pairs' within the entire array.

Learn more about Trie patterns.

Solution Approach

The solution involves constructing a Trie to handle the binary representations of numbers for quick retrieval and comparison. Here's a step-by-step explanation of how the solution is structured and how it works with the Trie:

Trie Data Structure

  • We define a Trie class to handle each bit of the 16-bit integers we are working with. Each node in the Trie has an array children of size 2 (one for bit 0 and another for bit 1) and a cnt variable to store the number of elements that follow the node's path.

Trie Insertion

  • The insert function takes an integer x and inserts it into the Trie bit by bit, starting from the most significant bit (15th bit). For each bit in x, if the corresponding children node doesn't exist, it creates a new Trie node. It then updates the current node's count.

Trie Searching

  • The search function calculates how many numbers in the Trie, when XORed with x, would produce a result lower than the given limit (which will be low or high + 1). It iterates through each bit (from the most significant to the least significant bit) and compares the bits of x and limit.

  • If the current bit of limit is 1, the function adds the count of the current bit path of x to the answer (because flipping the current bit of x could only increase the XOR result) and then proceeds to check the opposite path for further potential matches.

  • If the current bit of limit is 0, it means we cannot switch to the higher XOR value path, because doing so would lead to an XOR result higher than the limit. Therefore, it follows the path consistent with x's current bit.

The CountPairs Function

  • The countPairs function from the Solution class is where the algorithm starts applying the Trie structure. For every number x in the array nums, the function performs the following steps:

    1. Calls tree.search(x, high + 1) to find the number of elements that give an XOR with x less than or equal to high.

    2. Calls tree.search(x, low) to find the number of elements that give an XOR with x less than low.

    3. The difference of the two search results is the count of 'nice pairs' for the number x (since we want the XOR to be at least low and at most high).

  • After searching, the function then inserts x into the Trie to include it in the subsequent searches for the following numbers in nums.

  • Finally, the countPairs function returns the total number of 'nice pairs' found for the entire array nums.

The use of Trie to handle the binary representations of numbers and the clever bit manipulation during Trie traversal enables efficient calculation of 'nice pairs', which would be computationally intensive to obtain with a brute-force approach.

Example Walkthrough

Let’s use a small example to illustrate the solution approach detailed in the content above.

Suppose we have nums = [3, 10, 5, 25], low = 10, and high = 20. The binary representations of these numbers up to 5 bits for illustration are 0011, 1010, 0101, and 11001.

  • We first initialize our Trie and prepare to insert each number from the array and simultaneously search for 'nice pairs'.

  • Starting with the first number 3 (binary 0011), we insert it into the Trie. Since there are no previous numbers, the search for 'nice pairs' will return 0 for both high + 1 and low.

  • Next, consider the number 10 (binary 1010). We search the Trie for numbers that could form 'nice pairs':

    • For high condition (high + 1 = 21): The binary of 20 (limit) is 10101, and we traverse the Trie. We notice that the first bit in 1010 and 10101 is the same (bit 1) so we follow the same path (no flips). For the second bit, 0 in our number and 0 in high + 1, we proceed similarly. For the third bit, we have a 1 in high + 1 which means we could flip our current 1 to 0 to increase the XOR, but there is no such path (since we only have 3 in the Trie), so we keep the path. This process continues until we reach the end, and in this case, since 3 XOR 10 is 9 which is less than 10 (our low), there are no valid pairs for high so far.

    • For low condition: The binary of low is 01010, we traverse the Trie following a similar process but since 3 XOR 10 is lower than 10, this pair is not valid either.

  • The Trie now has 3 and 10. No 'nice pairs' have been found yet.

  • For the number 5 (binary 0101), we search the Trie for potential 'nice pairs':

    • For high condition (high + 1 = 21): During the Trie traversal, we find that 5 XOR 3 is 6, and 5 XOR 10 is 15. Both are within our range [10, 20], thus adding 2 to our 'nice pairs' count.

    • For low condition: Since our XOR results (6 and 15) are already greater than 10, we do not subtract any pairs for the low condition.

  • With 5 inserted, our Trie now contains 3, 10, and 5. We have found 2 'nice pairs' so far.

  • Finally, for number 25 (binary 11001), we search the Trie:

    • For high condition (high + 1 = 21): We find that 25 XOR 3 = 26, 25 XOR 10 is 27, and 25 XOR 5 is 28. All these results exceed the high value, so they are not valid.

    • For low condition: They also exceed low, so again, they are not subtracted from our 'nice pairs' count.

  • No additional 'nice pairs' are found, and 25 is inserted into the Trie.

The total number of 'nice pairs' found is 2: (5 XOR 3) and (5 XOR 10).

The Trie and search operations allow us to efficiently manage and calculate the XOR pairings without having to resort to a brute-force comparison, which would be less efficient for larger datasets.

Python Solution

1class TrieNode:
2    def __init__(self):
3        self.children = [None] * 2  # A binary trie, so 2 children, for 0 and 1
4        self.count = 0  # Maintain a count of numbers that pass through this node
5
6    def insert(self, number):
7        node = self
8        for i in range(15, -1, -1):  # Represent the number as a 16 bit integer
9            bit = (number >> i) & 1  # Extract the i-th bit of number
10            if node.children[bit] is None:
11                node.children[bit] = TrieNode()
12            node = node.children[bit]
13            node.count += 1  # Increment the count of numbers passing through this node
14          
15    # Search the trie to find the count of pairs with XOR less than limit
16    def search(self, number, limit):
17        node = self
18        result = 0
19        for i in range(15, -1, -1):  # Traverse from most significant bit to least
20            if node is None:
21                return result
22            bit = (number >> i) & 1
23            limit_bit = (limit >> i) & 1
24            if limit_bit:  # Checking if the bit is set in the limit.
25                # If so, all numbers with the same bit at this position
26                # will have a XOR less than limit.
27                if node.children[bit]:
28                    result += node.children[bit].count
29                # Move to the opposite bit to keep the XOR sum less than limit
30                node = node.children[bit ^ 1] 
31            else:
32                node = node.children[bit]
33        return result
34
35
36class Solution:
37    def countPairs(self, nums, low, high):
38        result = 0
39        trie = TrieNode()
40        for number in nums:
41            # Add count of valid pairs between number and numbers already in the trie
42            result += trie.search(number, high + 1) - trie.search(number, low)
43            trie.insert(number)
44        return result
45

Java Solution

1class Trie {
2    private Trie[] children = new Trie[2]; // Trie node children, one for bit 0, one for bit 1
3    private int count; // Number of elements that pass through this node
4
5    // Inserts a number into the trie
6    public void insert(int number) {
7        Trie node = this; // Start from the root
8        for (int i = 15; i >= 0; --i) { // Iterate over the bits of the number
9            int bit = (number >> i) & 1; // Extract the current bit
10            if (node.children[bit] == null) {
11                node.children[bit] = new Trie(); // Create new node if it doesn't exist
12            }
13            node = node.children[bit];
14            ++node.count; // Increment the count because we're adding a number
15        }
16    }
17
18    // Searches for numbers in the trie that have a XOR result with x below the given limit
19    public int search(int x, int limit) {
20        Trie node = this; // Start from the root
21        int answer = 0; // Initialize the answer
22        for (int i = 15; i >= 0 && node != null; --i) {
23            int bit = (x >> i) & 1; // Extract the current bit
24            if (((limit >> i) & 1) == 1) {
25                // If the bit is '1' in the limit, add count of numbers that have the opposite bit
26                if (node.children[bit] != null) {
27                    answer += node.children[bit].count;
28                }
29                node = node.children[bit ^ 1]; // Move to the opposite bit's child node if it exists
30            } else {
31                node = node.children[bit]; // If the bit is '0', just move to the same bit's child node
32            }
33        }
34        return answer;
35    }
36}
37
38class Solution {
39    // Counts the number of pairs (i, j) where i < j and XOR of nums[i] and nums[j] is in the range [low, high]
40    public int countPairs(int[] nums, int low, int high) {
41        Trie trie = new Trie(); // Initialize the trie
42        int answer = 0; // Initialize the number of valid pairs
43        for (int number : nums) {
44            // Search through the trie and check for valid pairs within bounds, high + 1 is used to include 'high' in range
45            answer += trie.search(number, high + 1) - trie.search(number, low); 
46            trie.insert(number); // Insert the number into the trie
47        }
48        return answer; // Return the final count of valid pairs
49    }
50}
51

C++ Solution

1#include <vector>
2using namespace std;
3
4class Trie {
5public:
6    Trie()
7        : children(2, nullptr)     // Initialize with 2 children as nullptrs for binary trie.
8        , count(0) {}
9
10    // Insert number x into Trie.
11    void insert(int x) {
12        Trie* node = this;
13        for (int i = 15; i >= 0; --i) { // Assuming 16-bit integer here.
14            int bit = (x >> i) & 1;
15            if (!node->children[bit]) {
16                node->children[bit] = new Trie();
17            }
18            node = node->children[bit];
19            ++node->count;  // Increment count at each node traversed.
20        }
21    }
22
23    // Search and count number of elements less than or equal to limit XOR'd with x.
24    int search(int x, int limit) {
25        Trie* node = this;
26        int ans = 0;
27        for (int i = 15; i >= 0 && node; --i) {
28            int bit = (x >> i) & 1;
29            int limitBit = (limit >> i) & 1;
30            if (limitBit) { // if current bit in limit is 1
31                if (node->children[bit]) {
32                    ans += node->children[bit]->count;
33                }
34                node = node->children[bit ^ 1]; // Move to the opposite bit node.
35            } else {
36                node = node->children[bit]; // Move to the same bit node as current bit of x
37            }
38        }
39        return ans;
40    }
41
42private:
43    vector<Trie*> children;  // Children is a vector of Trie pointers for the binary trie.
44    int count;  // Count of numbers in the Trie that have traversed through this node.
45};
46
47class Solution {
48public:
49    // Counts the number of pairs that have a bitwise XOR in [low, high].
50    int countPairs(vector<int>& nums, int low, int high) {
51        Trie* trie = new Trie(); // Create a new Trie.
52        int ans = 0;
53        for (int x : nums) {
54            // Count and subtract to get number of pairs with XOR in range.
55            ans += trie->search(x, high + 1) - trie->search(x, low);
56            trie->insert(x);  // Insert the number into the Trie.
57        }
58        return ans;  // Return the final answer.
59    }
60};
61

Typescript Solution

1type TrieNode = {
2    children: (TrieNode | null)[];
3    count: number;
4};
5
6// Initialize the Trie root node
7const root: TrieNode = { children: [null, null], count: 0 };
8
9// Insert number x into Trie
10function insert(x: number) {
11    let node = root;
12    for (let i = 15; i >= 0; --i) {
13        const bit = (x >> i) & 1;
14        if (!node.children[bit]) {
15            node.children[bit] = { children: [null, null], count: 0 };
16        }
17        node = node.children[bit]!;
18        ++node.count;
19    }
20}
21
22// Search and count number of elements less than or equal to limit XOR'd with x
23function search(x: number, limit: number) {
24    let node = root;
25    let ans = 0;
26    for (let i = 15; i >= 0; --i) {
27        const bit = (x >> i) & 1;
28        const limitBit = (limit >> i) & 1;
29        if (limitBit) {
30            if (node.children[bit]) {
31                ans += node.children[bit].count;
32            }
33            node = node.children[bit ^ 1];
34        } else {
35            node = node.children[bit];
36        }
37        if (!node) break;
38    }
39    return ans;
40}
41
42// Count the number of pairs that have a bitwise XOR in [low, high]
43function countPairs(nums: number[], low: number, high: number) {
44    let ans = 0;
45    for (const x of nums) {
46        ans += search(x, high + 1) - search(x, low);
47        insert(x);
48    }
49    return ans;
50}
51

Time and Space Complexity

Time Complexity

The time complexity of the insert and search methods in the Trie class is O(16) for each call which simplifies to O(1) since the loop runs for a fixed number of iterations (16 in this case), corresponding to the binary representation of the integers within the fixed range [0, 2^16 - 1].

The countPairs function calls insert once and search twice for every element in the nums list. If n is the number of elements in nums, then the overall time complexity of the countPairs function would be O(3n), which simplifies to O(n).

Hence, the total time complexity of the solution is O(n).

Space Complexity

The space complexity is determined by the size of the Trie data structure. In the worst case, if all the binary representations of the numbers in the list are unique, the trie could have up to 2^16 nodes (for a 16-bit integer). However, since the Trie is built as the numbers are added, and it only contains paths for numbers in the nums list, the space complexity would be O(16n) which simplifies to O(n) because each number in nums could theoretically end up contributing to 16 nodes in the Trie (one for each bit of the 16-bit integer).

Therefore, the total space complexity of the solution is O(n).

Learn more about how to find time and space complexity quickly.


Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫