Calculating Water

I came across this interesting programming challenge and couldn't help but dig in and solve it for myself. The solution below is a re-write of the approach I found in this book. My first (less efficient) solution involved a breadth first traversal of a histogram "graph" storing visited & filled cells in a HashSet. 

	/**
	* Trapping Water
	*
	* A one-dimensional container is specified by an array of
	* n nonnegative integers, specifying the eight of each
	* unit-width rectangle. Design an algorithm for computing
	* the capacity of the container.
	*
	* Ex: (X = Occupied, ▤ = Water)
	*
	* X X
	* XXX XXX
	* XXXX X XXXX▤▤▤X
	* X XXXX X X ==> X▤XXXX▤X▤X ==> 1+2+1+3 = 7
	* XXXXXXXXX X XXXXXXXXX▤X
	* 12134451203
	*/
	public class TrappingWater {
	/**
	* The key to this problem is realizing that the water
	* poured into the "container" will flow out of bounds
	* on the left and right sides (index < 0 &
	* index > array.length). If we are able to find the
	* highest point in the array (maxVal), we can divide
	* the array into two segments (left and right) around
	* this midpoint. For both left & right, there will be
	* a region of increasing altitude, and decreasing
	* altitude. In the regions of decreasing altitude,
	* this is where water will be stored. All we need to
	* do is find the local maxima in each left and right
	* region.
	*
	* Thus, the problem can be solved in 3 phases:
	* 1) Find the column index with the highest height
	* (maxIndex)
	* 2) Calculate water in the left hand side: For
	* i = [0, maxIndex-1] if we find a new local
	* maxima, store it, otherwise we have water
	* trapped and we can caluculate the volume
	* 3) Calculate water in the right hand side in a
	* similar manner for i = [maxIndex+1,
	* array.length-1]
	*
	* Solution is O(N) time, O(1) space
	*/
	private static int calculateWater(int[] heights) {
	// Arrays less than 3 can't contain water
	if (heights.length < 3 || heights == null) {
	return 0;
	}
	// 1) Calculate max height index
	int maxVal = 0, maxIndex = 0;
	for (int i = 0; i < heights.length; i++) {
	if (heights[i] > maxVal) {
	maxVal = heights[i];
	maxIndex = i;
	}
	}
	// 2) Calculate water in left portion of array
	int count = 0;
	int leftMax = 0;
	for (int i = 0; i < maxIndex - 1; i++) {
	if (heights[i] > leftMax) {
	leftMax = heights[i];
	} else {
	count += leftMax = heights[i];
	}
	}
	// 3) Calculate water in right portion of array
	int rightMax = 0;
	for (int i = heights.length - 1; i > maxIndex; i--) {
	if (heights[i] > rightMax) {
	rightMax = heights[i];
	} else {
	count += rightMax - heights[i];
	}
	}
	return count;
	}
	// --- Test code
	public static void main(String[] args) {
	int[] testData1 = {1};
	testCalculateWater(1, testData1, 0);
	int[] testData2 = {1, 2, 1, 3, 4, 4, 5, 1, 2, 0, 3};
	testCalculateWater(2, testData2, 7);
	int[] testData3 = {1, 2, 3, 4, 3, 2, 1};
	testCalculateWater(3, testData3, 0);
	int[] testData4 = {1, 0, 1, 0, 1, 0, 1, 0};
	testCalculateWater(4, testData4, 3);
	int[] testData5 = {Integer.MAX_VALUE,
	Integer.MAX_VALUE,
	Integer.MAX_VALUE,
	Integer.MAX_VALUE};
	testCalculateWater(5, testData5, 0);
	}
	private static void testCalculateWater(int testId,
	int[] input,
	int expectedVal) {
	int ret = calculateWater(input);
	if (ret == expectedVal) {
	System.out.println("PASS");
	} else {
	System.out.println(
	String.format(
	"Test %d: FAIL\n\tExpected return = %d\n\tActual return = %d",
	testId, expectedVal, ret)
	);
	}
	}
	}

view raw CalculateWater.java hosted with ❤ by GitHub

The follow up question to this is even more interesting:
"Solve this problem with an algorithm that accesses A's elements in order and can read an element only once. Use minimal additional space."