Edited code to complete assignment