FYI this algorithm is called the "Water filling algorithm" and is used extensively in Communications to optimize the allocation power for channels.

You can get a solution with simple Lagrangian method (which I believe is the linear solution).

http://www.eecs.berkeley.edu/~dtse/Chapters_PDF/Fundamentals...

(pages 183 - 185)