Improve rebalance