Code improvement based on homework 5