Small simplification in tryBalance

Instead of computing the bad/good node lists separately (in effect
iterating twice over the all_nodes list and also creating that list
twice from the container), let's use 'partition' and do a small
improvement.

Signed-off-by: Iustin Pop <iustin@google.com>
Reviewed-by: Guido Trotter <ultrotter@google.com>