An approach to writing bug-free Binary Search code

Pre-requisites

An idea about what binary search is, as this post is focused on how to write bug-free code for this.

Problem statement

There are multiple variants of binary search, I will cover these three:

Search the target element in a sorted array
Search the smallest element in a sorted array greater than the target element
Search the largest element in a sorted array lesser than the target element

Format of this post

In this post,

First, I will write standard code I have seen for these problems
Then, I will mention what problems I faced while using these standard codes and how I fell into bug traps,
Then, I will provide an approach how I write code such that I don't think twice about these general bugs.
Then some examples/applications.
Lastly, some caveats observed while using this approach.

Standard codes

The function will return the index as the answer instead of element, and -1 when there is no answer.

1. Search the target element in a sorted array

Standard code :

int search(vector<int> arr, int target) {
    int lo = 0, hi = arr.size() - 1;
	while(lo &lt;= hi) { // observe "&lt;=" operator here
	    int mid = lo + (hi-lo)/2;
		if(arr[mid] == target) return mid;
		if(arr[mid] &lt; target) lo = mid+1;
		else hi = mid-1;
	}
	return -1;
}
</int>

You may notice that I choose mid = lo + (hi-lo)/2 instead of mid = (lo + hi) / 2. That is because the latter one is prone to overflow. Consider the case when lo = 10 and hi = INT_MAX and check the output for both the expressions. 😊

2. Search the smallest element in a sorted array greater than the target element

A standard way to write this would be:

int greaterThan(vector<int> arr, int target) {
	int lo = 0, hi = arr.size() - 1;
	while(lo &lt; hi) { // observe this time it is "&lt;" operator, not "&lt;="
		int mid = lo + (hi-lo)/2;
		if(arr[mid] &lt;= target) lo = mid+1;
		else hi = mid;
	}
	return arr[lo] &gt; target ? lo : -1;
}
</int>

3. Search the largest element in a sorted array lesser than the target elementA standard way for this would be:

int smallerThan(vector<int> arr, int target) {
	int lo = 0, hi = arr.size() - 1;
	while(lo &lt; hi) {
		int mid = lo + (hi-lo+1)/2; // Note that I used ceil here otherwise it could go in infinite loop
		if(arr[mid] &lt; target) lo = mid;
		else hi = mid-1;
	}
	return arr[lo] &lt; target ? lo : -1;
}
</int>

There are few other implementations to solve the problem of infinite loop but using ceil looked cleaner way to me.

Introduction of bugs

As you can see, while writing Binary search (or its variants), you need to

think about the while loop condition,
infinite loop case,
choose between lo = mid vs lo = mid + 1,
choose between hi = mid vs hi = mid - 1

These are the main culprits to introduce bugs that I have experienced.

A bug-free approach

My approach to writing Binary-search related code is to:

Always keep an ans variable to store/update the answer. You have to initialize it correctly though. E.g. if you want to maximize something, start with the minimum that can be the answer (we will see this in the end of this post while solving a problem). In the end, we would be returning this ans variable without thinking twice.
Always make while loop condition to be while(lo <= hi), note the <= sign, and we will either make lo = mid+1 or hi = mid-1 (no more choosing here 😊), which would always avoid infinite loop condition, which is not the case when we do lo = mid or hi = mid and that introduces the infinite loop bugs.

That creates a pretty determinstic loop for us and we can play inside it.

1. Search the target element in a sorted array

int search(vector<int> arr, int target) {
    int lo = 0, hi = arr.size() - 1, ans = -1;
	while(lo &lt;= hi) {
		int mid = lo + (hi-lo)/2;
		if(arr[mid] == target) {
			ans = mid;
			break;
		} else if(arr[mid] &lt; target) lo = mid+1;
		else hi = mid-1;
	}
	return ans;
}
</int>

2. Search the smallest element in a sorted array greater than the target element

int greaterThan(vector<int> arr, int target) {
	int lo = 0, hi = arr.size() - 1, ans = -1;
	while(lo &lt;= hi) {
		int mid = lo + (hi-lo)/2;
		if(arr[mid] &gt; target) {
			ans = mid; // So far this could be the answer
			hi = mid-1; // But we can find better answer in left
		} else {
			lo = mid+1; // Try right
		}
	}
	return ans;
}
</int>

3. Search the largest element in a sorted array lesser than the target element

int smallerThan(vector<int> arr, int target) {
	int lo = 0, hi = arr.size() - 1, ans = -1;
	while(lo &lt;= hi) {
		int mid = lo + (hi-lo)/2; // Note that we don't have to think twice whether we should use ceil here or not, this will always work
		if(arr[mid] &lt; target) {
			ans = mid; // So far this could be the answer
			lo = mid+1; // But we can find better answer in right
		} else {
			hi = mid-1; // Try left
		}
	}
	return ans;
}
</int>

If you notice,

Last two codes are almost similar,
We didn't have to think about infinite loops,
We didn't have to think about while loop condition,
All we have to think about is the if-condition, and that would always depend on the problem statement.

Applications/Examples

1. Searching first and last occurrence of the target element in a sorted array

To find the first occurrence, whenever we find the element, we update our answer, and search for better answer in left (why?) Similarly for last occurrence, we update our answer whenever we find the element, and then search in right.

int firstOccurrence(vector<int> arr, int target) {
	int lo = 0, hi = arr.size() - 1, ans = -1; // these should always be initialized carefully
	while(lo &lt;= hi) { // this will never change
		int mid = lo + (hi-lo)/2; // this will never change
		if(arr[mid] == target) {
			ans = mid; // Update the answer
			hi = mid - 1; // Search in left for better answer
		} else if(arr[mid] &lt; target) lo = mid + 1;
		else hi = mid - 1;
	}
	return ans; // this will never change
}

int lastOccurrence(vector<int> arr, int target) {
	int lo = 0, hi = arr.size() - 1, ans = -1;
	while(lo &lt;= hi) {
		int mid = lo + (hi-lo)/2;
		if(arr[mid] == target) {
			ans = mid; // Update the answer
			lo = mid + 1; // Search in right for better answer
		} else if(arr[mid] &lt; target) lo = mid + 1;
		else hi = mid - 1;
	}
	return ans;
}
</int></int>

2. Binary search over monotonic functions

There are many problems which can be solved using binary search, not on some array, but on some monotonic function (a function that is either increasing or decreasing e.g. f(x) = x+1, as we increase x, the value of f(x) will always increase).

Let's solve 1011. Capacity To Ship Packages Within D Days to understand:

We can observe that if we keep the capacity very high, we can always ship, and with very low capacity, we can't ship, so we need to find the lowest possible ans such that it is shippable for capacity >= ans

int shipWithinDays(vector<int>&amp; weights, int D) {
	// Lets first define our feasible range
	int lo = weights[0];
	for(int w : weights) lo = max(lo, w); // we would at least need the capacity such that the largest weight can be shipped, so that would be lower bound
	int hi = 0;
	for(int w : weights) hi += w; // we would not need the capacity more than sum of all weights, so that would be upper bound
	int ans = hi; // since we need to minimize the capacity, start with hi, later we will update/minimize it
	while(lo &lt;= hi) { // No need to think here :D
		int mid = lo + (hi-lo)/2;
		if(isPossible(weights, mid, D)) {
			ans = mid; // Update answer
			hi = mid-1; // Try smaller capacities
		} else {
			lo = mid+1;
		}
	}
	return ans;
}

// Code to check whether it is possible to ship with this capacity
bool isPossible(vector<int>&amp; weights, int capacity, int D) {
	int cur = 0;
	for(int w : weights) {
		if(cur + w &lt;= capacity) {
			cur += w;
		} else {
			D--;
			if(D == 0) return false;
			cur = w;
		}
	}
	return true;
}
</int></int>

Another problem: 658. Find K Closest Elements

Please read lee215's explanation first. We start with some [left, right] range of indexes, from which we will find the index start, such that our answer will be arr[start:start+k]If it is better to have arr[mid:mid+k] compared to arr[mid+1:mid+k+1], that means we should update mid as answer and go in left direction i.e. right = mid-1And if we see it is better to have arr[mid+1:mid+k+1] compared to arr[mid:mid+k], that means we should not go with this mid, so just go in right direction i.e. left = mid+1

    def findClosestElements(self, arr: List[int], k: int, x: int) -&gt; List[int]:
        left, right = 0, len(arr)-k
        start = right // Initialization
        while left &lt;= right:
            mid = left + (right-left)//2
            if mid+k &lt; len(arr) and arr[mid+k]-x &lt; x-arr[mid]: // Note that `mid+k` can be out of bounds when mid=len(arr)-k
                left = mid + 1 // mid is not good start, it's better to move right and find the answer
            else:
                start = mid // mid is better than its right counterpart, so use it for now
                right = mid - 1 // but let's check if we can find better answer in left
        return arr[start:start+k]

Some caveats observed with this approach

Now that we have avoided infinite loop bug in binary search, doesn't mean no other bugs can come. Here are a few caveats we should consider:

Initialization of lo, hi, and ans is critical. We can't initialize lo to 0, and hi to INT_MAX blindly.
When accessing any other index than mid in array, check if it is in the bounds. For e.g. arr[mid+k], where k is some positive integer, may not be in the array bounds.

Tips and Tricks

If you are finding it hard to come up with a good inital value of ans, you can avoid it with the below idiom:

int greaterThan(vector<int> arr, int target) {
	int lo = 0, hi = arr.size() - 1, ans = -1, updated = false;
	while(lo &lt;= hi) {
		int mid = lo + (hi-lo)/2;
		if(arr[mid] &gt; target) {
			ans = mid; // So far this could be the answer
                        updated = true;
			hi = mid-1; // But we can find better answer in left
		} else {
			lo = mid+1; // Try right
		}
	}
	return updated ? -1 : ans;
}
</int>

Here we initialized one new variable updated to track whether ans is updated or not.

In the end, if ans is not updated (i.e. the variable updated is false), we return the default answer, else the variable ans.

Thanks for reaching to the end, hope this post was worth your time :)