Made slight improvment to algo