simplified code

list comprehensions are fun.
Merge pull request #171 from virgil/patch-3

simplified code