Cleaned up second exercise code