optimize performance of `balance_statespace`

[baggepinnen]

This function is being called in a lot of places so it makes sense to have it optimized, even if the code became a bit uglier

[JuliaControlBot]

This is an automated message. Plots were compared to references. 11/11 images have changed, see differences below. After pulling this PR, please update the reference images by creating a PR to ControlExamplePlots.jl here. Difference	Reference Image	New Image
:x: 0.047
:heavy_check_mark: 0.011
:heavy_check_mark: 0.001
:heavy_check_mark: 0.001
:heavy_check_mark: 0.007
:heavy_check_mark: 0.003
:heavy_check_mark: 0.0
:warning: 0.028
:x: 0.06
:heavy_check_mark: 0.0
:warning: 0.017

[codecov[bot]]

Codecov Report

Merging #852 (343c9b0) into master (5d4b138) will decrease coverage by 1.52%. The diff coverage is n/a.

@@            Coverage Diff             @@
##           master     #852      +/-   ##
==========================================
- Coverage   92.72%   91.21%   -1.52%     
==========================================
  Files           5        5              
  Lines         330      330              
==========================================
- Hits          306      301       -5     
- Misses         24       29       +5     

see 1 file with indirect coverage changes