Minor changes to solution