Counting the number of 1 bits in a sorted bit array 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 &lt;= 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)
        else:
            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
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 &lt; r:
        mid = l + (r - l) // 2
        if arr[mid] == 1:
            r = mid
        else:
            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.