Improve random number generator helper in the tree module

Taken from #3173

Ping @vene, @jnothman, @amueller
This indentation is a bit nonstandard. Otherwise, LGTM.