refactor the code
refactor the code
add problem 5