Fix balance function

This change fixes bug #5. It's my first pull request ever; please let me know if I've done anything wrong.