Work for data science homework