SumTree

Schaul et al. describe exactly this: "a 'sum-tree' data structure, where every node is the sum of its children, with the priorities as the leaf nodes, which can be efficiently updated and sampled from."

It is here rather than expressed with tensors because the obvious tensor formulation is a cumulative sum plus a search, and that is $O(n)$ in the buffer rather than $O(\log n)$ in it. Measured on this machine at a million leaves: the cumulative-sum form takes 196 ms to draw a batch of 32, the tree 0.20 ms — a factor of 973, against a training step of about 120 ms. The bottleneck is the algorithm, not the language, which is also why this is Python and not a new engine op: having fixed the complexity, a C++ tree would save a further 0.2% of a step.