Optimization for Problem 3

Adding a small optimization for Problem 3. Shouldn't affect run time by
much, but this is too obvious to not do.