Updated code to complete class assignment