Skip to Content »

Michi’s blog » Recursively counting numbers with fixed bit counts

 Recursively counting numbers with fixed bit counts

  • May 13th, 2012
  • 1:04 am

I ran across this problem in a reddit side-bar job-ad, and was intrigued by the task (description paraphrased to decrease googleability):

Write a function

uint64_t bitsquares(uint64_t a, uint64_t b);

such that it return the number of integers in [a,b] that have a square number of bits set to 1. Your function should run in less than O(b-a).

I think I see how to do it in something like logarithmic time. Here’s how:

First off, we notice that we can list all the squares between 0 and 64: these are 0, 1, 9, 16, 25, 36, 49, and 64. The function I will propose will run through a binary tree of depth 64, shortcutting through branches whenever it can. In fact; changing implementation language completely, I wonder if I cannot even write it comprehensively in Haskell.

The key insight I had was that whenever you try to find the number of numbers with a bitcount matching some element of some list within the bounds of 0b0000…0000 and 0b000…01111…11, then it reduces to a simple binomial coefficient — n choose k gives the number of numbers with k bits set among the n last. Furthermore, we can reduce the total size of the problem by removing a matching prefix from the two numbers we test from.

Hence, we trace how many bits off the top agree between the two numbers. We count the set bits among these, subtract them from each representative in the list of squares, giving us the counts we need to hit in the remainder.

Write a’ for a with the agreeing prefix removed, and similarly for b’. Then the total count is the count for the reduced things from a’ to 0b000…01…111 plus the count for the reduced things from 0 to b’. The reduction count for b’ needs to be 1 larger than the one for a’ since in one case, we are working with the prefix before the varying bit increases, and in the other, we work with the prefix after the varying bit increases — the latter count is not really from 0 to b’, but this is a useful proxy for the count from 0b0000…010…000 to b’ with the additional high bit set.

In code, I managed to boil this down to:

import Data.Word
import Data.Bits
import Data.List (elemIndices)

bitsquare :: Word64 -> Word64 -> Word64bitsquare a b = bitcountin a b squares -- # integers in [a,b] with square # of 1

squares = [1,4,9,16,25,36,49,64] :: [Word64]
allones = [fromIntegral (2^k - 1) | k <- [1..64]]

choose n 0 = 1
choose 0 k = 0
choose n k = (choose (n-1) (k-1)) * n `div` k

popCount :: Word64 -> Word64
popCount w = sum [1 | x <- [0..63], testBit w x]

-- # integers in [a,b] with 1-counts in counts
bitcountin :: Word64 -> Word64 -> [Word64] -> Word64
bitcountin a b counts
| a > b = 0
| a == b = if popCount b `elem` counts then 1 else 0 | (a == 0) && (b `elem` allones) = sum [choose n k | n <- [popCount b], k <- c
| otherwise = (bitcountin a' low [c-lobits | c <- counts, c>= lobits]) +
(bitcountin hi b' [c-hibits | c <- counts, c>= hibits])
agreements = [(testBit a n) == (testBit b n) | n <- [0..63]]
agreeI = elemIndices False agreements
prefixIndex = last agreeI
prefixCount = sum [1 | x <- [prefixIndex..63], testBit a x]
a' = a .&. (2^prefixIndex - 1)
b' = b .&. (2^prefixIndex - 1)
low = 2^prefixIndex - 1
hi = 0
lobits = prefixCount
hibits = prefixCount+1

5 People had this to say...

  • James Cook
  • May 14th, 2012
  • 17:49

For what it’s worth, it’s much simpler if you start off by defining a single-variable version of the problem (i.e., calculate bitSquareLTE n = bitsquare 0 n, and then just define the range version by subtraction.

An implementation of that approach is here:

  • Sam
  • May 14th, 2012
  • 18:00

I solved this in C instead. It is a bit wonky since you need to be careful to avoid integer overflow when you compute binomial coefficients (it’s better to just memoize Pascal’s triangle), but gcc’s __builtin_clzl() comes in handy for the actual puzzle. You can simplify the algorithm at the cost of slightly increased runtime by defining the function `bitCount 0 b counts’ instead and then defining the full function by subtraction.

  • Noah Easterly
  • May 14th, 2012
  • 18:41

popCount is already defined in Data.Bits.

  • Sam
  • May 14th, 2012
  • 19:41

P.S. I know that redefining bitCount to subtract is cheating, because the problem asks for an O(log(b-a)) solution. However, ‘bitCount a a k’ tells you how whether the population count of a is k or not; I don’t think that this can be done in constant time if k and a go to infinity.

  • Michi
  • May 15th, 2012
  • 18:38

@Noah: I know; however, it wasn’t in the version of the Haskell Platform I was working on. Hence my inclusion of the code.

Everyone else: Sure, I can see how the one-parameter version simplifies thinking about the problem. But it seems a bit wasteful to traverse the entire (underlying) bit decision tree when you only need to look at the span between the numbers given…

Want your say?

* Required fields. Your e-mail address will not be published on this site

You can use the following XHTML tags:
<a href="" title=""> <abbr title=""> <acronym title=""> <b> <blockquote cite=""> <cite> <code> <del datetime=""> <em> <i> <q cite=""> <strike> <strong>