See Gumbel-max trick at Tim Vieira's blog, and this follow-up post on Weighted reservoir sampling.

The Gumbel-max trick is widely used in deep learning from this paper: Categorical Reparameterization with Gumbel-Softmax. See also the Reparameterization Trick section 6.4 in Kim et al 2018 - A Tutorial on Deep Latent Variable Models of Natural Language