Clean up code for allocating teams

Shorten code
Avoid excessive len() calls
re-insert return statement
Merge pull request #1 from sakaru/master

Clean up code for allocating teams