Improve random number generator helper in the tree module

Taken from #3173

Ping @vene, @jnothman, @amueller