Adding problem3 code, but it's not finished yet