Updated code as per suggestions in review
2 files changed