recent Google interview puzzle on bitwise operation

This is a recent interview question from Google:

We define f(X, Y) as number of different corresponding bits in binary representation of X and Y. For example, f(2, 7) = 2, since binary representation of 2 and 7 are 010 and 111, respectively. The first and the third bit differ, so f(2, 7) = 2.

You are given an array of N positive integers, A1, A2 ,…, AN. Find sum of f(Ai, Aj) for all pairs (i, j) such that 1 ≤ i, j ≤ N

For example:

A=[1, 3, 5]

We return

f(1, 1) + f(1, 3) + f(1, 5) + f(3, 1) + f(3, 3) + f(3, 5) + f(5, 1) + f(5, 3) + f(5, 5) =

0 + 1 + 1 + 1 + 0 + 2 + 1 + 2 + 0 = 8

I could think of this solution which is O(n^2)

int numSetBits(unsigned int A) {
    int count  = 0;

    while(A != 0) {
        A = A & (A-1);
        count++;
    }

    return count;
}

int count_diff_bits(int a, int b)
{
    int x = a ^ b;

    return numSetBits(x);
}

for (i = 0; i < n; i++)
   for (j = 0; j < n; j++) {
       sum += count_diff_bits(A[i], A[j]);
   }
}

Another approach i can think of is (considering that each element contains only one binary digit):

Start from the end of the array
keep a count of 1's and 0's found so far
If the current element is 1, then it will contribute count_of_zeros to the final sum
Continue like this till we reach the start of the array.

Is this approach correct.

Solution

Iterate the array, and count number of "on" bits in each bit index, for example [1, 3, 5]:

Now, for each bit counter, calculate:

[bit count] * [array size - bit count] * 2

and sum for all bits...

With example above:

3 * (3 - 3) * 2 = 0
1 * (3 - 1) * 2 = 4
1 * (3 - 1) * 2 = 4
          total = 8

To show why this works, lets look at a subset of the problem, using a single bit. Let's see what happens if we have an array with: [1, 1, 0, 0, 1, 0, 1]. Our count is 4 and size is 7. If we examine the first bit with all the bits in the array (including self, as in the question), we get:

1 xor 1 = 0
1 xor 1 = 0
1 xor 0 = 1
1 xor 0 = 1
1 xor 1 = 0
1 xor 0 = 1
1 xor 1 = 0

As can be seen, the contribution of this bit is the number of "off" bits. The same holds true for any other "on" bit. We could say that each "on" bit counts as the number of "off" bits:

[bit count] * [array size - bit count]

And where does the multiplication by 2 comes from? well, since we do the same with the "off" bits, except that for these, the contribution is the number of "on" bits:

[array size - bit count] * [bit count]

which of course is the same as above, and we can just multiply...

Complexity is O(n*k) where k is number of bits (32 in your code).