A bit of pythonic approach.
Merge pull request #2 from akash0x53/master

code beatufication and pythonization thanks to @akash0x53