Binary Search in a 2D Matrix
In our problem last week, we covered a complex problem that used a binary search. Today, we’ll apply binary search again to solidify our understanding of it. This time, instead of extra algorithmic complexity, we’ll start adding some data structure complexity. We’ll be working with a 2D Matrix instead of basic arrays.
To learn more about data structures and algorithms in Haskell, you should take a look at our Solve.hs course! In particular, you’ll cover multi-dimensional arrays in module 2, and you’ll learn how to write algorithms in Haskell in module 3!
The Problem
Today’s problem is Search a 2D Matrix, and the description is straightforward. We’re given a 2D m x n matrix, as well as a target number. We have to return a boolean for whether or not that number is in the Matrix.
This is trivial with a simple scan, but we have an additional constraint that lets us solve the problem faster. The matrix is essentially ordered. Each row is non-decreasing, and the first element of each successive row is no smaller than the last element of the preceding row.
This allows us to get a solution that is O(log(n + m)), a considerable improvement over a linear scan.
The Algorithm
The algorithm is simple as well. We’ll do two binary searches. First, we’ll search over the rows to identify the last row which could contain the element. Then we’ll do a binary search of that row to see if the element is present or not.
We’ll have a slightly different form to our searches compared to last time. In last week’s problem, we knew we had to find a valid index for our search. Now, we may find that no valid index exists.
So we’ll structure our search interval in a semi-open fashion. The first index in our search interval is inclusive, meaning that it could still be a valid index. The second index is exclusive, meaning it is the lowest index that we consider invalid.
In mathematical notation, we would represent such an interval with a square bracket on the left and a parenthesis on the right. So if that interval is [0, 4), then 0, 1, 2, 3 are valid values. The interval [2,2) would be considered empty, with no valid values. We’ll see how we apply this idea in practice.
Rust Solution
We don’t have that many terms to define at the start of this solution. We’ll save the size of both dimensions, and then prepare ourselves for the first binary search by assigning low as 0 (the first potential “valid” answer), hi as m (the lowest “invalid” answer), and creating our output rowWithTarget value. For this, we also assign m, an invalid value. If we fail to re-assign rowWithTarget in our binary search, we want it assigned to an easily testable invalid value.
pub fn search_matrix(matrix: Vec<Vec<i32>>, target: i32) -> bool {
let m = matrix.len();
let n = matrix[0].len();
let mut low = 0;
let mut hi = m;
let mut rowWithTarget = m;
...
}Now we write our first binary search, looking for a row that could contain our target value. We maintain the typical pattern of binary search, using the loop while (low < hi) and assigning mid = (low + hi) / 2.
pub fn search_matrix(matrix: Vec<Vec<i32>>, target: i32) -> bool {
...
while (low < hi) {
let mid: usize = (low + hi) / 2;
if (matrix[mid][0] > target) {
hi = mid;
} else if (matrix[mid][n - 1] < target) {
low = mid + 1;
} else {
rowWithTarget = mid;
break;
}
}
if (rowWithTarget >= m) {
return false;
}
...
}If the first element of the row is too large, we know that mid is “invalid”, so we can assign it as hi and continue. If the last element is too small, then we reassign low as mid + 1, as we want low to still be a potentially valid value.
Otherwise, we have found a potential row, so we assign rowWithTarget and break. If, after this search, rowWithTarget has the “invalid” value of m, we can return false, as there are no valid values.
Now we just do the same thing over again, but within rowWithTarget! We reassign low and hi (as n this time) to reset the while loop. And now our comparisons will look at the specific value matrix[rowWithTarget][mid].
pub fn search_matrix(matrix: Vec<Vec<i32>>, target: i32) -> bool {
...
low = 0;
hi = n;
while (low < hi) {
let mid: usize = (low + hi) / 2;
if (matrix[rowWithTarget][mid] > target) {
hi = mid;
} else if (matrix[rowWithTarget][mid] < target) {
low = mid + 1;
} else {
return true;
}
}
return false;
}Again, we follow the same pattern of re-assigning low and hi. If we don’t hit the return true case in the loop, we’ll end up with return false at the end, because we haven’t found the target.
Here’s the full solution:
pub fn search_matrix(matrix: Vec<Vec<i32>>, target: i32) -> bool {
let m = matrix.len();
let n = matrix[0].len();
let mut low = 0;
let mut hi = m;
let mut rowWithTarget = m;
while (low < hi) {
let mid: usize = (low + hi) / 2;
if (matrix[mid][0] > target) {
hi = mid;
} else if (matrix[mid][n - 1] < target) {
low = mid + 1;
} else {
rowWithTarget = mid;
break;
}
}
if (rowWithTarget >= m) {
return false;
}
low = 0;
hi = n;
while (low < hi) {
let mid: usize = (low + hi) / 2;
if (matrix[rowWithTarget][mid] > target) {
hi = mid;
} else if (matrix[rowWithTarget][mid] < target) {
low = mid + 1;
} else {
return true;
}
}
return false;
}Haskell Solution
In our Haskell solution, the main difference of course will be using recursion for the binary search. However, we’ll also change up the data structure a bit. In the Rust framing of the problem, we had a vector of vectors of values. We could do this in Haskell, but we could also use Array (Int, Int) Int. This lets us map row/column pairs to numbers in a more intuitive way.
import qualified Data.Array as A
search2DMatrix :: A.Array (Int, Int) Int -> Int -> Bool
search2DMatrix matrix target = ...
where
((minR, minC), (maxR, maxC)) = A.bounds matrixAnother unique feature of arrays is that the bounds don’t have to start from 0. We can have totally custom bounding dimensions for our rows and columns. So instead of using m and n, we’ll need to use the min and max of the row and column dimensions.
So now let’s define our first binary search, looking for the valid row. As we did last week, the input to our function will be two Int values, for the low and hi. As in our Rust solution we’ll access the first and last element of the row defined by the “middle” of low and hi, and compare them against the target. We make recursive calls to searchRow if the row isn’t valid.
search2DMatrix :: A.Array (Int, Int) Int -> Int -> Bool
search2DMatrix matrix target = result
where
((minR, minC), (maxR, maxC)) = A.bounds matrix
searchRow :: (Int, Int) -> Int
searchRow (low, hi) = if low >= hi then maxR + 1 else
let mid = (low + hi) `quot` 2
firstInRow = matrix A.! (mid, minC)
lastInRow = matrix A.! (mid, maxC)
in if firstInRow > target
then searchRow (low, mid)
else if lastInRow < target
then searchRow (mid + 1, hi)
else mid
rowWithTarget = searchRow (minR, maxR + 1)
result = rowWithTarget <= maxR && ...Instead of m, we have maxR + 1, which we use as the initial hi value, as well as a return value in the base case where low meets hi. We can return a result of False if rowWithTarget does not come back with a value smaller than maxR.
Now for our second search, we follow the same pattern, but now we’re returning a boolean. The base case returns False, and we return True if we find the value in rowWithTarget at position mid. Here’s what that looks like:
search2DMatrix :: A.Array (Int, Int) Int -> Int -> Bool
search2DMatrix matrix target = result
where
...
rowWithTarget = searchRow (minR, maxR + 1)
searchCol :: (Int, Int) -> Bool
searchCol (low, hi) = low < hi &&
let mid = (low + hi) `quot` 2
val = matrix A.! (rowWithTarget, mid)
in if val > target
then searchCol (low, mid)
else if val < target
then searchCol (mid + 1, hi)
else True
result = rowWithTarget <= maxR && searchCol (minC, maxC + 1)You’ll see we now use the outcome of searchCol for result. And this completes our solution! Here’s the full code:
search2DMatrix :: A.Array (Int, Int) Int -> Int -> Bool
search2DMatrix matrix target = result
where
((minR, minC), (maxR, maxC)) = A.bounds matrix
searchRow :: (Int, Int) -> Int
searchRow (low, hi) = if low >= hi then maxR + 1 else
let mid = (low + hi) `quot` 2
firstInRow = matrix A.! (mid, minC)
lastInRow = matrix A.! (mid, maxC)
in if firstInRow > target
then searchRow (low, mid)
else if lastInRow < target
then searchRow (mid + 1, hi)
else mid
rowWithTarget = searchRow (minR, maxR + 1)
searchCol :: (Int, Int) -> Bool
searchCol (low, hi) = low < hi &&
let mid = (low + hi) `quot` 2
val = matrix A.! (rowWithTarget, mid)
in if val > target
then searchCol (low, mid)
else if val < target
then searchCol (mid + 1, hi)
else True
result = rowWithTarget <= maxR && searchCol (minC, maxC + 1)Conclusion
Next week, we’ll stay on the subject of 2D matrices, but we’ll learn about array mutation. This is a very tricky subject in Haskell, so make sure to come back for that article!
To learn how these data structures work in Haskell, read about Solve.hs, our Haskell Data Structures & Algorithms course!