How to select a binary tree node uniformly at random

Question

The exercise I'm trying to solve is

You are implementing a binary search tree class from scratch, which, in addition, to insert, find and delete, has a method getRandomNode() which returns a random node from the tree. All nodes should be equally likely to be chosen. Design and implement an algorithm for getRandomNode(), and explain how you would implement the rest of the methods.

The answer from the book is:

 1  class TreeNode {
 2    private int data;
 3    public TreeNode left;
 4    public TreeNode right;
 5    private int size = 0;
 ...
12    public TreeNode getRandomNode() {
13      int leftSize = left == null ? 0 : left.size();
14      Random random = new Random();
15      int index = random.nextInt(size);
16      if (index < leftSize) {
17        return left.getRandomNode();
18      } else if (index == leftSize) {
19        return this;
20      } else {
21        return right.getRandomNode();
22      }
23    }
...
55  }

But here is the problem. With this algorithm, I don't see how nodes are equally likely to be chosen. In line 16, it says if leftsize > index, where index is a number from 0 to size, then the algorithm will continue with the left node, otherwise the right node. It only works when the tree has a depth of 2. When the tree is taller, the probability of each node being chosen will not be equal.

Am I wrong? Does this algorithm work?

Answer 1

The algorithm works just fine.

Note that each node's size field tells you the total number of nodes in the subtree rooted at that node. Throughout this answer, I'm just going to assume that a child's size is zero if the child is null. Thus, for any node in the tree, we have

this.size == left.size + 1 + right.size

Consider the root of the tree. If we choose a node of the tree uniformly at random, then it should be in the left subtree with probability left.size/size, it should be the root node with probability 1/size and it should be in the right subtree with probability right.size/size. By the equation above, these three probabilities add up to 1. If you look at lines 15–22, the algorithm implements exactly this probability distribution. It generates a uniform random integer in the range 0 .. size-1. If this number is in the range 0 .. left.size-1, the algorithm picks a random node from the left subtree; if it's equal to left.size, it picks the root and, otherwise, it picks a random node from the right subtree.

And the argument I gave doesn't depend on the node you consider being the root, so the recursive calls behave as they should. The algorithm doesn't care about the depth of the nodes: it just recurses down the subtree and chooses subtrees with the correct probability.

Answer 2

I want to show that this algorithm indeed randomly chooses a node according to the uniform distribution.

Consider the root node. We have two subtrees: the left subtree with the size $s_l$ and the right subtree with the size $s_r$. The total number of nodes in the tree is $size = 1 + s_r + s_l$. Now, I want to prove that distribution is uniform, i.e. for any node $e$, $P(e)=\frac{1}{size}$ (probability that the algorithm chooses $e$).

There are three possibilities, either $e$ is the root, or $e$ is in the left subtree, or $e$ is the right subtree.

If $e$ is the root node, then index = random.nextInt(size)uniformly selects an integer between $0$ and $size-1$, and the statement if (index == leftSize) is true with the probability $\frac{1}{size}$.

If the node $e$ is in the left subtree then the probability that the statement if (index < leftSize) is true, $P(InLeft)$, is equal to $\frac{s_l}{size}$ and the conditional probability that the node $e$ is chosen given it is in the left subtree is $P(e \mid InLeft)=\frac{1}{s_l}$. Thus the probability [$e$ is in the left subtree] AND [$e$ is chosen] is equal to $$P(InLeft \cap e \text{ is chosen }) = P(InLeft)P(e \mid InLeft) = \frac{s_l}{size}\frac{1}{s_l} = \frac{1}{size}$$

If the node $e$ is in the right subtree then using the same argument as in the previous case we have $P(InRight)P(e\mid InRight) = \frac{s_r}{size}\frac{1}{s_r} = \frac{1}{size}$.

Thus the distribution is uniform.

Answer 3

Consider an arbitrary node $n_k$ which is reached from the root $n_0$ along path $(n_0, n_1, \dots, n_k)$. We will explicitly compute the probability of $n_k$ being chosen.

Observe that when visiting $n_{i-1}$ with $i>0$, that node itself is selected with probability

$\qquad\displaystyle \frac{1}{\operatorname{size}(n_i)}$, and

child $n_i$ is recursed to with probability

$\qquad\displaystyle \frac{\operatorname{size}(n_i)}{\operatorname{size}(n_{i-1})}$.

Hence, because all the choices are stochastically independent, the probability to select $n_k$ is

$\qquad\displaystyle \frac{\operatorname{size}(n_1)}{\operatorname{size}(n_0)} \cdot \frac{\operatorname{size}(n_2)}{\operatorname{size}(n_1)} \cdot \dots \cdot \frac{\operatorname{size}(n_k)}{\operatorname{size}(n_{k-1})} \cdot \frac{1}{\operatorname{size}(n_k)} \quad=\quad \frac{1}{\operatorname{size}(n_0)}$,

what was to show.

Fill in the details for a formal proof.