Added rebalance method and tests, they pass