Few small changes

Basically just cleaning it up so that it can work on a bunch of
different data sets. BTW, don't Fork this if you're intending to pass
this coursework off as your own. The Uni frowns on that.