Code modified for assignment
Minor changes