Search code examples
algorithmbit-manipulationxor

Number of pair in array that satisfy Bitwise Equation

I found this problem in a contest. The question is:

You are given an array of N non-negative integers (A1, A2, A3,..., An) and an integer M. Your task is to find the number of unordered pairs of array elements (X,Y) that satisfies the following bitwise equation:

2 * set_bits(X|Y) = M + set_bits(X ⊕ Y)

Note:

  • set_bits(n) represents the number of ones in the binary represenataion of an integer n.
  • X|Y represents the bitwise OR of integer X and Y.
  • X ⊕ Y represents the bitwise XOR of integer X and Y.
  • The unordered pair of array elements is pair (Ai, Aj) where 1 ≤ i < j ≤ N.

Print the number of unordered pairs of array elements that satisfy the above bitwise equation.

  1. Sample Input 1:
    N=4 M=2
    arr = [3, 0, 4, 5]
    Sample Output: 2
    2 pairs are (3,0) and (0,5)

  2. Sample Input 2:
    N=8 M=2
    arr = [3, 0, 4, 5, 6, 8, 1, 8]
    Sample Output: 9

Is there any other way except brute force to solve this equation?

Solution

  • A solution with time complexity O(n) exists if the time complexity of set_bits is O(1).

    First, we are going to rephrase the condition (the bitwise equation) a bit. Assume a pair of elements (X, Y) is given. Let c_01 represent the number of digits where X is 0 but Y is 1, c_10 represent the number digits where X is 1 and Y is 0, and c_11 represent the number of digits where X and Y are 1. For example, when X=5 and Y=1 (X=101, Y=001), c_01 = 0, c_10 = 1, c_11 = 1. Now, the condition can be rewritten as

    2 * (c_01 + c_10 + c_11) = M + (c_01 + c_10)

    because set_bits(X|Y) is equal to c_01 + c_10 + c_11 and set_bits(X^Y) is equal to c_01 + c_10. We can reorder the equation into

    c_01 + c_10 + 2*c_11 = M

    by moving the term on the right to the left side. Now, realize that set_bits(X) = c_10 + c_11. Applying this information on the equation we get

    c_01 + c_11 = M - set_bits(X)

    Now, also realize that set_bits(Y) = c_01 + c_11. The equation becomes

    set_bits(Y) = M - set_bits(X)

    or

    set_bits(X) + set_bits(Y) = M

    The problem has turned into counting the number of pairs such that the number of set bits in the first element plus the number of set bits in the second element is equal to M. This can be done in linear time assuming you can compute set_bits in constant time.