I wrote a function in python to count the number of 1-bits in a sorted bit array. I'm basically using binary search but the code seems unnecessarily long and awkward. I did test it for several different cases and got the correct output but am looking to write a cleaner code if possible.

Note: I want to use a binary-search inspired solution that theoretically takes O(lg n ) time.

def countOnes(arr):
    l, r = 0, len(arr)-1
    while l <= r:
        mid = l + (r-l)/2
        if mid == 0 and arr[mid] == 1:
            return len(arr)
        if mid == len(arr)-1 and arr[mid] == 0:
            return 0
        if arr[mid] == 1 and arr[mid-1] == 0:
            return len(arr) - mid
        elif arr[mid] == 1 and arr[mid-1] == 1:
            r = mid - 1
        elif arr[mid] == 0 and arr[mid+1] == 1:
            return len(arr) - (mid+1)
            l = mid + 1

print countOnes([0, 0, 0]) # got 0
print countOnes([0, 0, 1, 1]) # got 2
print countOnes([1, 1, 1, 1]) # got 4
print countOnes([0, 1, 1, 1]) # got 3
print countONes([1, 1, 1, 1, 1]) # got 5
  • \$\begingroup\$ Where do these sorted bit-arrays come from? Would it be possible to represent them as pairs (number of zeros, number of ones)? \$\endgroup\$ – Gareth Rees Jan 23 '18 at 13:30

I agree with your assessment of the code being unnecessarily long and awkward. It looks like the root cause is an opportunistic optimization.

The bad news is that most of the secondary tests (like arr[mid - 1] == 0 and arr[mid - 1] == 1) are bound to fail, so not only they contribute to awkwardness of the code - they also hurt the performance (by about a factor of 2, if I am not mistaken).

I recommend to keep it simple, and test only what is necessary:

    l, r = 0, len(arr)
    while l < r:
        mid = l + (r - l) // 2
        if arr[mid] == 1:
            r = mid
            l = mid + 1
    return len(arr) - l

Notice that working on a semi-open range (that is, mid is beyond the range) also makes it cleaner.

| improve this answer | |
  • \$\begingroup\$ I see that makes a lot of sense. Can you explain one detail? Why is r updated to mid and not mid-1 but l is updated to mid+1 and not just mid. \$\endgroup\$ – grayQuant Jan 23 '18 at 0:38
  • 1
    \$\begingroup\$ @grayQuant We are maintaining semi-openness. We want guarantee that arr[r] (where it is legal) to be 1. r = mid - 1 may break the invariant. \$\endgroup\$ – vnp Jan 23 '18 at 0:41
  • \$\begingroup\$ @grayQuant To clarify, an invariant or the loop is that there are no zeroes starting from r . r = mid maintains it. \$\endgroup\$ – vnp Jan 23 '18 at 0:52
  • \$\begingroup\$ This is the best answer because it actually takes my code and makes it much cleaner and readable. I did post it on codereview after all. I'm not looking for other approaches to do this trivial task. \$\endgroup\$ – grayQuant Jan 24 '18 at 7:07

This removes the learning part of this (if your goal is to try to learn to implement a binary search in this context), but I contend the most pythonic solution would be to use bisect.bisect_right. It gives you the rightmost insertion point for a value via binary search. The rightmost insertion point for a 0 would be the number of zeros, so you just subtract this from the length to get the number of ones:

from bisect import bisect_right

def count_ones(bits):
    return len(bits) - bisect_right(bits, 0)

Some notes:

  • Use PEP8 style, snake_case function names
  • Add a """Docstring."""
  • I think bits is a more descriptive name than arr
  • If counting is such a core operation, you should wrap this array in a class that counts the number as you modify the array
  • An array is a pretty sparsely packed representation of a bitfield. The bitarray package is a more efficient C implementation (spacewise). What's fantastic about it is that it conforms to the list API, so you should in theory be able to just pass a bitarray into count_ones instead of a list and it will just work™
  • You could potentially squeeze more performance out of this if needed. By representing it as a bytearray, you'd want to find the first non-zero byte (which can be achieved by bisect_right(bits, 0)) and then check the next one to see how many bits into it the first 1 is. This has the advantage of avoiding the shifting and anding to extract specific bits at each point considered by the binary search. The downside here is you loose the ability to address individual bits using a list like interface (and doing the | and << in python may be slower than bitarrays C implementation). But that all said, you'd definitely want to benchmark; this is just a guess.

Edit: Dug a little bit into that last point. Turns out you can do .tobytes() which returns a reference to the backing bytes of the bitarray. So, you could do the bisect_right on it as described above (for what I would imagine would be a faster binary search), but still have the advantage of treating the DS like a list of bits (with operations at the speed the C implementation).

| improve this answer | |

You can't get much simpler than:


This has the benefit of working even if the array isn't sorted.

| improve this answer | |
  • 2
    \$\begingroup\$ Sorry this is an exercise in binary search where I want O(lg(n)) time and not O(n) which the sum function would be. \$\endgroup\$ – grayQuant Jan 23 '18 at 0:07
  • \$\begingroup\$ How big are these arrays? For small enough n, O(n) can actually be faster. \$\endgroup\$ – Scott Hunter Jan 23 '18 at 14:07
  • \$\begingroup\$ Doesn't matter, this is not a practical application. Just an exercise in binary search \$\endgroup\$ – grayQuant Jan 24 '18 at 7:04

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.