Two Pointer Algorithms
We’re now on to part 5 of our series comparing Haskell and Rust solutions for LeetCodeproblems. You can also look at the previous parts (Part 1, Part 2, Part 3, Part 4) to get some more context on what we’ve learned so far comparing these two languages.
For a full look at problem solving in Haskell, check out Solve.hs, our latest course! You’ll get full breakdowns on the processes for solving problems in Haskell, from basic list and loop problems to advanced algorithms!
The Problem
Today we’ll be looking at a problem called Trapping Rain Water. In this problem, we’re given a vector of heights, which form a sort of 1-dimensional topology. Our job is to figure out how many units of water could be collected within the topology.
As a very simple example, the input [1,0,2]
could collect 1 unit of water. Here’s a visualization of that system, where x
shows the topology and o
shows water we collect:
x
xox
We can never collect any water over the left or right “edges” of the array, since it would flow off. The middle index of our array though is lower than its neighbors. So we take the lower of these neighboring values, and we see that we can collect 1
unit of water in this system.
For a bigger example that collects water, we might have the input [4, 2, 1, 1, 3, 5]
. Here’s what that looks like:
x
x o o o o x
x o o o x x
x x o o x x
x x x x x x
The total water here is 9.
A flat system like [2,2,2]
, or a system that looks like a peak [1,2,3,2,1]
cannot collect any water, so we should return 0 in these cases.
The Algorithm
There are a couple ways to solve this. One approach would be a two-pass solution, similar to what we used in Product of Array Except Self. We loop from the left side, tracking the maximum water we can store in each unit based on its left neighbors. Then we loop again from the right side and compare the maximum we can store based on the right neighbors to the prior value from the left. This solution is O(n)
time, but O(n)
space as well.
A more optimal solution for this problem is a two-pointer approach that can use O(1)
additional space. In this kind of solution, we look at the left and right of the input simultaneously. Each step of the way, we make a decision to either increase the “left pointer” or decrease the “right pointer” until they meet in the middle. Each time we move, we get more information about our solution.
In this particular problem, we’ll track the maximum value we’ve seen from the left side and the maximum value we’ve seen from the right side. As we traverse each index, we update both sides for the current left and right indices if we have a new maximum.
The crucial step is to see that if the current “left max” is smaller than the current “right max”, we know how much water can be stored at the left index. This is just the left max minus the left index. Then we can increment the left index.
If the opposite is true, we calculate how much water can be stored at the right index, and decrease the right index.
So we keep a running tally of these sums, and we end our loop when they meet in the middle.
Rust Solution
We can describe our algorithm as a simple while loop. This loop goes until the left index exceeds the right index. The loop needs to track 5 values:
- Left Index
- Right Index
- Left Max
- Right Max
- Total sum so far
So let’s write the setup portion of the loop:
pub fn trap(height: Vec<i32>) -> i32 {
let mut leftMax = -1;
let mut rightMax = -1;
let mut leftI = 0;
let mut rightI = height.len() - 1;
let mut total = 0;
while leftI <= rightI {
...
}
}
A subtle thing…the constraints on the LeetCode problem are that the length is at least 1. But to handle length 0 cases, we would need a special case. Rust uses unsigned integers for vector length, so taking height.len() - 1
on a length-0 vector would give the maximum integer, and this would mess up our loop and indexing.
Within the while
loop, we run the algorithm.
- Adjust
leftMax
andrightMax
if necessary. - If
leftMax
is not larger, recurse, incrementingleftI
and adding tototal
from the left - If
rightMax
is smaller, decrementrightI
and add total from the right
And at the end, we return our total
!
pub fn trap(height: Vec<i32>) -> i32 {
let n = height.len();
if n <= 1 {
return 0;
}
let mut leftMax = -1;
let mut rightMax = -1;
let mut leftI = 0;
let mut rightI = n - 1;
let mut total = 0;
while leftI <= rightI {
// Step 1
leftMax = std::cmp::max(leftMax, height[leftI]);
rightMax = std::cmp::max(rightMax, height[rightI]);
if leftMax <= rightMax {
// Step 2
total += leftMax - height[leftI];
leftI += 1;
} else {
// Step 3
total += rightMax - height[rightI];
rightI -= 1;
}
}
return total;
}
Haskell Solution
Now that we’ve seen our Rust solution with a single loop, let’s remember our process for translating this idea to Haskell. With a two-pointer loop, the way in which we traverse the elements of the input is unpredictable, thus we need a raw recursive function, rather than a fold or a map.
Since we’re tracking 5 integer values, we’ll want to write a loop
function that looks like this:
-- (leftIndex, rightIndex, leftMax, rightMax, sum)
loop :: (Int, Int, Int, Int, Int) -> Int
Knowing this, we can already “start from the end” and figure out how to invoke our loop from the start of our function:
trapWater :: V.Vector Int -> Int
trapWater input = loop (0, n - 1, -1, -1, 0)
where
n = V.length input
loop :: (Int, Int, Int, Int, Int) -> Int
loop = undefined
In writing our recursive loop, we’ll start with the base case. Once leftI
is the bigger index, we return the total.
trapWater :: V.Vector Int -> Int
trapWater input = loop (0, n - 1, -1, -1, 0)
where
n = V.length input
loop :: (Int, Int, Int, Int, Int) -> Int
loop (leftI, rightI, leftMax, rightMax, total) = if leftI > rightI then total
else …
Within the else
case, we just follow our algorithm, with the same 3 steps we saw with Rust.
trapWater :: V.Vector Int -> Int
trapWater input = loop (0, n - 1, -1, -1, 0)
where
n = V.length input
-- (leftIndex, rightIndex, leftMax, rightMax, sum)
loop :: (Int, Int, Int, Int, Int) -> Int
loop (leftI, rightI, leftMax, rightMax, total) = if leftI > rightI then total
else
-- Step 1
let leftMax' = max leftMax (input V.! leftI)
rightMax' = max rightMax (input V.! rightI)
in if leftMax' <= rightMax'
-- Step 2
then loop (leftI + 1, rightI, leftMax', rightMax', total + leftMax' - input V.! leftI)
-- Step 3
else loop (leftI, rightI - 1, leftMax', rightMax', total + rightMax' - input V.! rightI)
And we have our Haskell solution!
Conclusion
If you’ve been following this whole series so far, hopefully you’re starting to get a feel for comparing basic algorithms in Haskell and Rust (standing as a proxy for most loop-based languages). In general, we can write loops as recursive functions in Haskell, capturing the “state” of the list as the input parameter for that function.
In particular cases where each iteration deals with exactly one element of an input list, we can employ folds as a tool to simplify our functions. But the two-pointer algorithm we explored today falls into the general recursive category.
To learn the details of understanding these problem solving techniques, take a look at our course, Solve.hs! You’ll learn everything from basic loop and list techniques, to advanced data structures and algorithms!